"""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/", 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"})