memory-system/auth/api_auth.py

"""API Key + HMAC Signature Authentication for Memory System

Auth flow:
  1. Client sends X-API-Key, X-Timestamp, X-Signature headers
  2. X-Signature = HMAC-SHA256(api_secret, timestamp + request_body)
  3. Server validates key exists, is active, timestamp within 5min, signature matches
  4. Authenticated team_id is injected into flask.g for downstream use
"""
import hashlib
import hmac
import json
import logging
import os
import secrets
import time
from datetime import datetime
from functools import wraps

import pymysql
from flask import Blueprint, current_app, g, request

logger = logging.getLogger(__name__)

# Timestamp tolerance in seconds (5 minutes)
TIMESTAMP_TOLERANCE = 300

api_auth_bp = Blueprint("admin", __name__, url_prefix="/admin")


# Helpers

def _get_db():
    ms = current_app.memory_system
    return ms.mysql._get_conn()


def _verify_signature(api_secret, timestamp, body):
    message = timestamp.encode() + body
    return hmac.new(api_secret.encode(), message, hashlib.sha256).hexdigest()


def _lookup_api_key(api_key):
    conn = _get_db()
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cur:
            cur.execute(
                "SELECT id, team_id, api_key, secret, name, is_active "
                "FROM api_keys WHERE api_key = %s",
                (api_key,),
            )
            return cur.fetchone()
    finally:
        conn.close()


def _touch_api_key(key_id):
    conn = _get_db()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "UPDATE api_keys SET last_used = NOW() WHERE id = %s",
                (key_id,),
            )
            conn.commit()
    finally:
        conn.close()


# Auth middleware

def require_auth():
    path = request.path

    # Skip health check, root landing page, and admin endpoints
    if path in ("/", "/health") or path.startswith("/admin"):
        return None

    api_key = request.headers.get("X-API-Key")
    timestamp = request.headers.get("X-Timestamp")
    signature = request.headers.get("X-Signature")

    if not api_key or not timestamp or not signature:
        return current_app.json_response(
            {"ok": False, "error": "Missing authentication headers. Required: X-API-Key, X-Timestamp, X-Signature",
             "code": "AUTH_MISSING_HEADERS"}, 401)

    try:
        ts = float(timestamp)
    except (ValueError, TypeError):
        return current_app.json_response(
            {"ok": False, "error": "Invalid timestamp format",
             "code": "AUTH_INVALID_TIMESTAMP"}, 401)

    now = time.time()
    if abs(now - ts) > TIMESTAMP_TOLERANCE:
        return current_app.json_response(
            {"ok": False, "error": "Timestamp expired (5 min tolerance)",
             "code": "AUTH_TIMESTAMP_EXPIRED"}, 401)

    key_record = _lookup_api_key(api_key)
    if not key_record:
        return current_app.json_response(
            {"ok": False, "error": "Invalid API key",
             "code": "AUTH_INVALID_KEY"}, 401)

    if not key_record["is_active"]:
        return current_app.json_response(
            {"ok": False, "error": "API key is disabled",
             "code": "AUTH_KEY_DISABLED"}, 401)

    body = request.get_data()
    expected = _verify_signature(key_record["secret"], timestamp, body)
    if not hmac.compare_digest(signature, expected):
        return current_app.json_response(
            {"ok": False, "error": "Invalid signature",
             "code": "AUTH_INVALID_SIGNATURE"}, 401)

    g.team_id = key_record["team_id"]
    g.api_key_id = key_record["id"]
    _touch_api_key(key_record["id"])
    return None


# Admin endpoints (localhost only)

@api_auth_bp.before_request
def _restrict_admin_to_localhost():
    remote = request.remote_addr or ""
    if remote not in ("127.0.0.1", "::1", "localhost"):
        return current_app.json_response(
            {"ok": False, "error": "Admin endpoints are localhost-only",
             "code": "FORBIDDEN_ADMIN"}, 403)


@api_auth_bp.route("/api-keys", methods=["POST"])
def create_api_key():
    d = request.json or {}
    team_id = d.get("team_id")
    name = d.get("name", "")
    if not team_id:
        return current_app.json_response(
            {"ok": False, "error": "team_id required",
             "code": "VALIDATION_ERROR"}, 400)

    if not current_app.memory_system.get_team(team_id):
        return current_app.json_response(
            {"ok": False, "error": f"Team {team_id} not found",
             "code": "TEAM_NOT_FOUND"}, 404)

    key_id = secrets.token_hex(16)
    api_key = "msk_" + secrets.token_hex(24)
    secret = secrets.token_hex(32)

    conn = _get_db()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO api_keys (id, team_id, api_key, secret, name, is_active) "
                "VALUES (%s, %s, %s, %s, %s, TRUE)",
                (key_id, team_id, api_key, secret, name),
            )
            conn.commit()
    finally:
        conn.close()

    return current_app.json_response({
        "ok": True,
        "data": {
            "id": key_id, "team_id": team_id, "api_key": api_key,
            "secret": secret, "name": name,
            "message": "Store the secret securely. It will NOT be shown again.",
        },
    }, 201)


@api_auth_bp.route("/api-keys", methods=["GET"])
def list_api_keys():
    conn = _get_db()
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cur:
            cur.execute(
                "SELECT id, team_id, api_key, name, is_active, created_at, last_used "
                "FROM api_keys ORDER BY created_at DESC")
            rows = cur.fetchall()
    finally:
        conn.close()

    for row in rows:
        for k, v in row.items():
            if isinstance(v, datetime):
                row[k] = v.isoformat()

    return current_app.json_response({"ok": True, "data": rows})


@api_auth_bp.route("/api-keys/<key_id>", methods=["DELETE"])
def disable_api_key(key_id):
    conn = _get_db()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "UPDATE api_keys SET is_active = FALSE WHERE id = %s",
                (key_id,))
            conn.commit()
            affected = cur.rowcount
    finally:
        conn.close()

    if affected == 0:
        return current_app.json_response(
            {"ok": False, "error": "API key not found",
             "code": "NOT_FOUND"}, 404)

    return current_app.json_response({"ok": True, "message": "API key disabled"})