app.py

"""
抖音港版登录页恢复 + 账号同步 - FastAPI 后端 (v2.1)

v2.1 变化:
  ✅ 保留全部原有接口:
       /api/verify
       /api/submit_task       (v1 恢复, 带 log)
       /api/v2/submit_task    (v2 恢复, 仅 progress + done)
       /api/admin/*           (管理员)
  ✅ 新增接口:
       /api/v2/sync           mobile→standard 账号同步 (SSE, 仅 progress + done)

安全约定:
  1. SSH 密码 / ADB token 绝对不落盘、不落库、不写日志
  2. 请求处理完立即从内存释放 (Python 自然回收)
  3. 日志只记: 时间、来源 IP、卡密前 6 位、步骤名、成功/失败; 不记 raw_input
"""

import asyncio
import json
import logging
import os
import queue
import secrets
import sqlite3
import d1db  # 可插拔 DB 后端 (SQLite 本地 / D1 云)
import sys
import threading
import time
from contextlib import contextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException, Header, Request, Body
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette.sse import EventSourceResponse

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import restore
import sync    # v2.1 新增
import hubble  # v2.6 新增: mobile → hubble 同步

# ----------------- 配置 -----------------

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "data")
STATIC_DIR = os.path.join(BASE_DIR, "static")
LOG_DIR = os.path.join(BASE_DIR, "logs")
DB_PATH = os.path.join(DATA_DIR, "keys.db")
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(LOG_DIR, exist_ok=True)

ADMIN_PASSWORD = os.environ.get("HKDY_ADMIN_PASSWORD", "admin")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(LOG_DIR, "app.log"), encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
log = logging.getLogger("hkdy")


# ----------------- DB -----------------

DEFAULT_USES = 9999


def init_db():
    con = d1db.connect(DB_PATH)
    con.execute("""
        CREATE TABLE IF NOT EXISTS keys (
            key         TEXT PRIMARY KEY,
            status      TEXT NOT NULL DEFAULT 'unused',
            uses_left   INTEGER NOT NULL DEFAULT 9999,
            created_at  INTEGER NOT NULL,
            used_at     INTEGER
        )
    """)
    try:
        con.execute(f"ALTER TABLE keys ADD COLUMN uses_left INTEGER NOT NULL DEFAULT {DEFAULT_USES}")
    except sqlite3.OperationalError:
        pass
    con.commit()
    con.close()


@contextmanager
def db():
    con = d1db.connect(DB_PATH)
    con.row_factory = sqlite3.Row
    try:
        yield con
    finally:
        con.close()


def new_key():
    a = secrets.token_hex(2).upper()
    b = secrets.token_hex(2).upper()
    return f"VIP-{a}-{b}"


def check_auth(authz: Optional[str]) -> Optional[str]:
    """返回角色: 'admin' / 'user' / None."""
    if not authz:
        return None
    token = authz.strip()
    if token == "ADMIN_TOKEN":
        return "admin"
    with db() as c:
        row = c.execute("SELECT * FROM keys WHERE key=?", (token,)).fetchone()
        if row:
            return "user"
    return None


def mask(s: str, keep: int = 6) -> str:
    if not s:
        return ""
    if len(s) <= keep:
        return "*" * len(s)
    return s[:keep] + "*" * (len(s) - keep)


def _deduct_uses_if_user(role, token):
    """普通用户扣次数, 返回扣后 uses_left; 管理员返回 None."""
    if role != "user":
        return None
    with db() as c:
        row = c.execute("SELECT uses_left FROM keys WHERE key=?", (token,)).fetchone()
        if not row or row["uses_left"] <= 0:
            raise HTTPException(402, "卡密使用次数已用尽")
        uses_left = row["uses_left"] - 1
        c.execute(
            "UPDATE keys SET uses_left=?, status=?, used_at=? WHERE key=?",
            (uses_left, "used", int(time.time()), token),
        )
        c.commit()
    return uses_left


# ----------------- App -----------------

app = FastAPI(title="HK-DYBack API v2.1")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)

init_db()


@app.post("/api/verify")
async def verify(request: Request, body: dict = Body(...)):
    key = (body.get("key") or "").strip()
    ip = request.client.host if request.client else "?"

    if not key:
        raise HTTPException(400, "卡密不能为空")

    if key == ADMIN_PASSWORD:
        log.info(f"verify admin from {ip}")
        return {"role": "admin", "token": "ADMIN_TOKEN"}

    with db() as c:
        row = c.execute("SELECT * FROM keys WHERE key=?", (key,)).fetchone()
        if not row:
            log.info(f"verify fail key={mask(key)} from {ip}")
            raise HTTPException(401, "无效卡密")
        if row["uses_left"] <= 0:
            log.info(f"verify exhausted key={mask(key)} from {ip}")
            raise HTTPException(402, "卡密使用次数已用尽")
    log.info(f"verify ok key={mask(key)} uses_left={row['uses_left']} from {ip}")
    return {"role": "user", "token": key, "uses_left": row["uses_left"]}


# --------- v1 恢复接口 (保留, 带完整日志推送) ---------

@app.post("/api/submit_task")
async def submit_task(request: Request, body: dict = Body(...),
                      authorization: Optional[str] = Header(None)):
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")

    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")

    ip = request.client.host if request.client else "?"
    token = authorization.strip()

    uses_left_after = _deduct_uses_if_user(role, token)
    log.info(f"task start role={role} key={mask(token)} uses_left={uses_left_after} from {ip}")

    q: "queue.Queue[dict]" = queue.Queue()
    done_event = threading.Event()
    result_box: dict = {}

    def logger(msg):
        q.put({"type": "log", "msg": str(msg)})

    def progress(val):
        q.put({"type": "progress", "value": int(val)})

    def worker():
        try:
            ok, png, bk = restore.run_restore(raw, DATA_DIR, log=logger, progress=progress)
            result_box.update({"ok": bool(ok), "backup": bk, "uses_left": uses_left_after})
        except Exception as e:
            logger(f"❌ 错误: {e}")
            result_box.update({"ok": False, "error": str(e), "uses_left": uses_left_after})
        finally:
            done_event.set()

    threading.Thread(target=worker, daemon=True).start()

    async def sse_gen():
        loop = asyncio.get_event_loop()
        while True:
            try:
                ev = await loop.run_in_executor(None, q.get, True, 1.0)
            except queue.Empty:
                if done_event.is_set() and q.empty():
                    break
                continue
            yield {"data": json.dumps(ev, ensure_ascii=False)}
            if ev.get("type") == "done":
                break
            if done_event.is_set() and q.empty():
                yield {"data": json.dumps(
                    {"type": "done", **result_box}, ensure_ascii=False
                )}
                break
        log.info(f"task end role={role} key={mask(token)} ok={result_box.get('ok')}")

    return EventSourceResponse(sse_gen())


# --------- v2 精简接口 (只推 progress / done) ---------

def _run_sse_task(runner, raw, scratch_dir, role, token, ip, uses_left_after, tag):
    """通用 SSE 包装: 跑 runner(raw, scratch_dir, log, progress), 推 progress/done.

    runner 签名需与 restore.run_restore / sync.run_sync 对齐:
        (raw_input_text, scratch_dir, log, progress) -> (ok, png, bk)
    """
    q: "queue.Queue[dict]" = queue.Queue()
    done_event = threading.Event()
    result_box: dict = {"uses_left": uses_left_after}

    def logger(msg):
        # v2 不推送文本日志, 仅服务端记录
        pass

    def progress(val):
        q.put({"type": "progress", "value": int(val)})

    def worker():
        try:
            ok, _png, _bk = runner(raw, scratch_dir, log=logger, progress=progress)
            result_box.update({"ok": bool(ok)})
        except Exception as e:
            result_box.update({"ok": False, "error": str(e)})
        finally:
            done_event.set()
            q.put({"type": "done", **result_box})

    threading.Thread(target=worker, daemon=True).start()

    async def sse_gen():
        loop = asyncio.get_event_loop()
        last_val = -1
        while True:
            try:
                ev = await loop.run_in_executor(None, q.get, True, 5.0)
            except queue.Empty:
                if done_event.is_set() and q.empty():
                    break
                continue
            if ev.get("type") == "progress":
                if ev["value"] == last_val:
                    continue
                last_val = ev["value"]
            yield {"data": json.dumps(ev, ensure_ascii=False)}
            if ev.get("type") == "done":
                break
        log.info(f"{tag} end role={role} key={mask(token)} ok={result_box.get('ok')}")

    return EventSourceResponse(sse_gen())


@app.post("/api/v2/submit_task")
async def submit_task_v2(request: Request, body: dict = Body(...),
                         authorization: Optional[str] = Header(None)):
    """v2 港版登录页恢复 (原有, 行为不变)"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")

    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")

    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)

    log.info(f"v2 task start role={role} key={mask(token)} uses_left={uses_left_after} from {ip}")
    return _run_sse_task(restore.run_restore, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 task")


@app.post("/api/v2/sync")
async def submit_sync_v2(request: Request, body: dict = Body(...),
                          authorization: Optional[str] = Header(None)):
    """v2.1 新增: mobile→standard 账号同步 (Plan A, 同 v2 SSE 协议)"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")

    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")

    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)

    log.info(f"v2 sync start role={role} key={mask(token)} uses_left={uses_left_after} from {ip}")
    return _run_sse_task(sync.run_sync, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 sync")


# --------- v2.6 Hubble 接口 ---------

@app.post("/api/v2/hubble")
async def submit_hubble_v2(request: Request, body: dict = Body(...),
                            authorization: Optional[str] = Header(None)):
    """mobile→hubble 同步+启动"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")
    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")
    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)
    log.info(f"v2 hubble start role={role} key={mask(token)} from {ip}")
    return _run_sse_task(hubble.run_hubble, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 hubble")


@app.post("/api/v2/hubble_sync")
async def submit_hubble_sync_v2(request: Request, body: dict = Body(...),
                                 authorization: Optional[str] = Header(None)):
    """仅 mobile→hubble 数据同步"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")
    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")
    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)
    log.info(f"v2 hubble_sync start role={role} key={mask(token)} from {ip}")
    return _run_sse_task(hubble.run_hubble_sync, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 hubble_sync")


@app.post("/api/v2/hubble_launch")
async def submit_hubble_launch_v2(request: Request, body: dict = Body(...),
                                   authorization: Optional[str] = Header(None)):
    """仅 hubble 启动页面"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")
    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")
    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)
    log.info(f"v2 hubble_launch start role={role} key={mask(token)} from {ip}")
    return _run_sse_task(hubble.run_hubble_launch, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 hubble_launch")


@app.post("/api/v2/hubble_quick")
async def submit_hubble_quick_v2(request: Request, body: dict = Body(...),
                                  authorization: Optional[str] = Header(None)):
    """快速模式: 只发 intent, 不等不验证"""
    role = check_auth(authorization)
    if not role:
        raise HTTPException(401, "未授权")
    raw = body.get("raw_input") or ""
    if not raw or len(raw) < 30:
        raise HTTPException(400, "输入缺失或过短")
    ip = request.client.host if request.client else "?"
    token = authorization.strip()
    uses_left_after = _deduct_uses_if_user(role, token)
    log.info(f"v2 hubble_quick start role={role} key={mask(token)} from {ip}")
    return _run_sse_task(hubble.run_hubble_quick, raw, DATA_DIR,
                         role, token, ip, uses_left_after, tag="v2 hubble_quick")


# --------- 管理员接口 (保留) ---------

@app.get("/api/admin/list_keys")
def list_keys(authorization: Optional[str] = Header(None)):
    if check_auth(authorization) != "admin":
        raise HTTPException(403, "需要管理员")
    with db() as c:
        rows = c.execute(
            "SELECT key,status,uses_left,created_at,used_at FROM keys ORDER BY created_at DESC"
        ).fetchall()
    return [dict(r) for r in rows]


@app.post("/api/admin/set_uses")
def set_uses(body: dict = Body(...),
             authorization: Optional[str] = Header(None)):
    if check_auth(authorization) != "admin":
        raise HTTPException(403, "需要管理员")
    k = (body.get("key") or "").strip()
    n = int(body.get("uses", DEFAULT_USES))
    if n < 0:
        n = 0
    with db() as c:
        c.execute("UPDATE keys SET uses_left=? WHERE key=?", (n, k))
        c.commit()
    return {"ok": True, "uses_left": n}


@app.post("/api/admin/generate_keys")
def gen_keys(body: dict = Body(...),
             authorization: Optional[str] = Header(None)):
    if check_auth(authorization) != "admin":
        raise HTTPException(403, "需要管理员")
    n = max(1, min(int(body.get("count", 1)), 50))
    uses = int(body.get("uses", DEFAULT_USES))
    if uses < 1:
        uses = DEFAULT_USES
    new = []
    now = int(time.time())
    with db() as c:
        for _ in range(n):
            for _try in range(5):
                k = new_key()
                try:
                    c.execute(
                        "INSERT INTO keys(key, status, uses_left, created_at) VALUES(?,?,?,?)",
                        (k, "unused", uses, now),
                    )
                    new.append(k)
                    break
                except sqlite3.IntegrityError:
                    continue
        c.commit()
    return {"keys": new, "uses_left": uses}


@app.post("/api/admin/delete_key")
def del_key(body: dict = Body(...),
            authorization: Optional[str] = Header(None)):
    if check_auth(authorization) != "admin":
        raise HTTPException(403, "需要管理员")
    k = (body.get("key") or "").strip()
    with db() as c:
        c.execute("DELETE FROM keys WHERE key=?", (k,))
        c.commit()
    return {"ok": True}


@app.get("/api/health")
def health():
    return {"status": "ok", "version": "2.6"}


# --------- 静态 ---------
# 仅当 static/ 存在且非空时挂载; API-only 部署不放 static/ 就跳过
if os.path.isdir(STATIC_DIR) and any(os.scandir(STATIC_DIR)):
    app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")