226 lines
7.5 KiB
Python
226 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
register_token.py — 将用户的 hc_session token 注册到 AI Token 表，
|
||
供 AI 直接调用后端 API。
|
||
|
||
用户只需提供一次，后端存储后 AI 后续就能用这个 token 访问。
|
||
|
||
用法：
|
||
python3 register_token.py
|
||
python3 register_token.py <token值> # 静默模式（用于自动化）
|
||
"""
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import secrets
|
||
import sys
|
||
import webbrowser
|
||
from datetime import datetime, timedelta
|
||
from pathlib import Path
|
||
from urllib.parse import urlencode
|
||
|
||
# ---------- 项目路径 ----------
|
||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||
BACKEND_DIR = SCRIPT_DIR.parent
|
||
sys.path.insert(0, str(BACKEND_DIR))
|
||
|
||
from config import Config
|
||
from database import connection, now_iso
|
||
|
||
|
||
def _sha256(value: str) -> str:
|
||
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def _normalize_user_from_token(token: str) -> dict | None:
|
||
"""用 token 查找用户，不存在则返回 None"""
|
||
session_hash = _sha256(token)
|
||
current = now_iso()
|
||
with connection() as conn:
|
||
row = conn.execute(
|
||
"""SELECT s.id AS session_id,u.* FROM auth_sessions s
|
||
JOIN auth_users u ON u.id=s.user_id
|
||
WHERE s.session_hash=? AND s.revoked_at IS NULL AND s.expires_at>?""",
|
||
(session_hash, current),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
conn.execute(
|
||
"UPDATE auth_sessions SET last_seen_at=? WHERE id=?",
|
||
(current, row["session_id"]),
|
||
)
|
||
return {
|
||
"id": str(row["ding_user_id"]),
|
||
"name": row["name"],
|
||
"ding_user_id": row["ding_user_id"],
|
||
}
|
||
|
||
|
||
def ensure_ai_tokens_table() -> None:
|
||
"""建 AI Token 表（如果不存在）"""
|
||
with connection() as conn:
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS ai_access_tokens (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
token_hash TEXT NOT NULL UNIQUE,
|
||
user_id TEXT NOT NULL,
|
||
user_name TEXT NOT NULL,
|
||
description TEXT DEFAULT '',
|
||
created_at TEXT NOT NULL,
|
||
expires_at TEXT NOT NULL,
|
||
revoked_at TEXT,
|
||
last_used_at TEXT,
|
||
use_count INTEGER NOT NULL DEFAULT 0
|
||
)
|
||
""")
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_ai_token_hash ON ai_access_tokens(token_hash)"
|
||
)
|
||
|
||
|
||
def register_ai_token(token: str, user_info: dict, description: str = "") -> dict:
|
||
"""注册 / 更新 AI 访问 token"""
|
||
token_hash = _sha256(token)
|
||
now = now_iso()
|
||
expires = (datetime.now().astimezone() + timedelta(days=365)).isoformat(
|
||
timespec="seconds"
|
||
)
|
||
with connection() as conn:
|
||
existing = conn.execute(
|
||
"SELECT id FROM ai_access_tokens WHERE token_hash=?",
|
||
(token_hash,),
|
||
).fetchone()
|
||
if existing:
|
||
conn.execute(
|
||
"""UPDATE ai_access_tokens
|
||
SET last_used_at=?, use_count=use_count+1
|
||
WHERE token_hash=?""",
|
||
(now, token_hash),
|
||
)
|
||
action = "updated"
|
||
else:
|
||
conn.execute(
|
||
"""INSERT INTO ai_access_tokens
|
||
(token_hash,user_id,user_name,description,created_at,expires_at,last_used_at,use_count)
|
||
VALUES (?,?,?,?,?,?,?,1)""",
|
||
(
|
||
token_hash, user_info["ding_user_id"],
|
||
user_info["name"], description, now, expires, now,
|
||
),
|
||
)
|
||
action = "registered"
|
||
return {"action": action, "user": user_info["name"], "expires_at": expires}
|
||
|
||
|
||
def revoke_ai_token(token: str) -> bool:
|
||
token_hash = _sha256(token)
|
||
with connection() as conn:
|
||
conn.execute(
|
||
"UPDATE ai_access_tokens SET revoked_at=? WHERE token_hash=? AND revoked_at IS NULL",
|
||
(now_iso(), token_hash),
|
||
)
|
||
affected = conn.total_changes
|
||
return affected > 0
|
||
|
||
|
||
def list_ai_tokens() -> list[dict]:
|
||
with connection() as conn:
|
||
rows = conn.execute(
|
||
"""SELECT user_name,description,created_at,expires_at,last_used_at,use_count,revoked_at
|
||
FROM ai_access_tokens WHERE revoked_at IS NULL
|
||
ORDER BY last_used_at DESC"""
|
||
).fetchall()
|
||
return [dict(r) for r in rows]
|
||
|
||
|
||
# ---------- 主入口：交互式注册流程 ----------
|
||
def main() -> None:
|
||
ensure_ai_tokens_table()
|
||
|
||
backend = os.environ.get("HC_BACKEND_URL", f"http://localhost:{Config.PORT}").rstrip("/")
|
||
|
||
print()
|
||
print("═" * 55)
|
||
print(" 🔐 Holy Crab · AI Token 注册")
|
||
print("═" * 55)
|
||
print()
|
||
print(" 本工具用于将你的 hc_session token 注册给 AI 访问权限。")
|
||
print()
|
||
print(" Step 1：请从浏览器获取你的 hc_session token")
|
||
print()
|
||
print(" 操作步骤：")
|
||
print(" 1. 在浏览器中打开 Holy Crab 前端并完成钉钉登录")
|
||
print(" 2. 打开浏览器 DevTools (F12) → Application → Cookies")
|
||
print(" 3. 找到你的后端地址，点击展开 Cookie 列表")
|
||
print(" 4. 复制 'hc_session' 的 Value（很长的一串字符）")
|
||
print()
|
||
|
||
# 检查是否传入参数
|
||
if len(sys.argv) > 1:
|
||
token = sys.argv[1].strip()
|
||
if token in ("--help", "-h"):
|
||
print("用法: python3 register_token.py [token]")
|
||
sys.exit(0)
|
||
print(f" 检测到命令行传入 token: {token[:8]}...{token[-4:]}")
|
||
else:
|
||
print(" 直接回车打开浏览器说明网页（推荐）：")
|
||
print()
|
||
print(" 或者手动输入 token（粘贴后回车）：")
|
||
token = input(" > ").strip()
|
||
|
||
if not token:
|
||
# 打开说明网页
|
||
help_url = (
|
||
f"{backend}/api/v1/auth/dingtalk/login"
|
||
f"?redirect={urlencode({'r': '/__help__'})}"
|
||
)
|
||
print()
|
||
print(f" 正在打开说明网页：{help_url}")
|
||
webbrowser.open(help_url)
|
||
print()
|
||
print(" 请按上述步骤获取 token 后再次运行：")
|
||
print(f" python3 {sys.argv[0]} <your_token>")
|
||
sys.exit(0)
|
||
|
||
print()
|
||
print(" 🔍 正在验证 token...")
|
||
user = _normalize_user_from_token(token)
|
||
if not user:
|
||
print()
|
||
print(" ❌ token 无效或已过期")
|
||
print(" 请确认：")
|
||
print(" 1. token 是否正确（注意不要有前后空格）")
|
||
print(" 2. 是否已通过钉钉登录")
|
||
print(" 3. token 是否已过期（默认14天）")
|
||
sys.exit(1)
|
||
|
||
print(f" ✅ 验证成功！用户：{user['name']}")
|
||
print()
|
||
result = register_ai_token(
|
||
token, user,
|
||
description=f"AI助手访问 (注册IP/来源: register_script)"
|
||
)
|
||
print(f" ✅ AI Token 注册成功")
|
||
print(f" 用户：{result['user']}")
|
||
print(f" 有效期至：{result['expires_at']}")
|
||
print()
|
||
print(" 📋 AI 后续使用方式：")
|
||
print()
|
||
print(" 方式 A — 设置环境变量（推荐，每次会话有效）：")
|
||
print(f" export HC_SESSION='{token}'")
|
||
print()
|
||
print(" 方式 B — 写入文件（持久化）：")
|
||
print(f" echo '{token}' > ~/.holy_crab_cookie")
|
||
print()
|
||
print(" 方式 C — 写入配置（自动读取）：")
|
||
print(f" echo 'HC_SESSION={token}' >> ~/.holy_crab_env")
|
||
print(" # scripts 会自动从 ~/.holy_crab_env 读取环境变量")
|
||
print()
|
||
print("═" * 55)
|
||
print()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|