初始化：记忆系统源代码上传（已脱敏）
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html） - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
This commit is contained in:
zqf
@@ -0,0 +1,2 @@
|
||||
"""Authentication module for Memory System API"""
|
||||
from auth.api_auth import require_auth, api_auth_bp
|
||||
@@ -0,0 +1,214 @@
|
||||
"""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"})
|
||||
Reference in New Issue
Block a user