holy后端
This commit is contained in:
225
scripts/register_token.py
Normal file
225
scripts/register_token.py
Normal file
@@ -0,0 +1,225 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user