shirobana-bot/identity.py

# identity.py — 统一身份层（平台层，与 channel.py / tg_channel.py / qq_channel.py 同级）
#
# 目标：每个使用者只有一个与平台无关的身份 id（uid）。TG 账号首次对话自动开卡；
#       QQ 账号通过「QQ 侧发起 + TG 侧确认」绑定到同一个 uid。业务状态一律以 uid 为键，
#       QQ 侧不再存在任何独立配置；未绑定的 QQ 用户只能看到绑定指引。
#
# 存储（目录由 configure() 注入，本模块不感知调用方的目录结构）：
#   users/<uid>.json   权威记录，每用户一个文件
#   index.json         生成的映射表（tg / qq → uid），供人查看与快速加载
#   pending.json       QQ 绑定待确认请求
#
# id 段位：QQ 号历史上限约 4.29e9（uint32）；TG 用户 ns = 1e10 + raw；TG 群 ns = 1e12 + |chat_id|。
#          故 [4.3e9, 1e10) 无任何业务 id 占用：uid 取 UID_BASE 起递增，GID_BASE 预留给第二步的
#          群身份（本模块暂不使用），二者错开，日志里也不会与 QQ/TG id 混淆。
import json
import os
import random
import threading
import time
from datetime import datetime, timedelta, timezone

UID_BASE = 9000000001
GID_BASE = 8000000001
PENDING_TTL = 600          # 绑定确认码有效期（秒）
MIN_QQ_LEN = 5
MAX_QQ_LEN = 11
TOKEN_LEN = 6
TZ8 = timezone(timedelta(hours=8))

_lock = threading.RLock()
_base_dir = ""
_admin_tg_raw = set()
_users = {}
_tg_index = {}
_qq_index = {}
_pending = {}
_next_seq = 1


# ==================== 基础工具 ====================
def _stamp():
    return datetime.now(TZ8).strftime("%Y-%m-%d %H:%M:%S")


def _tg_offset():
    """TG 用户命名空间偏移（与 tg_channel.TG_USER_OFFSET 同一来源）"""
    from tg_channel import TG_USER_OFFSET
    return TG_USER_OFFSET


def _users_dir():
    return os.path.join(_base_dir, "users")


def _user_path(uid):
    return os.path.join(_users_dir(), "%d.json" % int(uid))


def _index_path():
    return os.path.join(_base_dir, "index.json")


def _pending_path():
    return os.path.join(_base_dir, "pending.json")


def _read_json(path, default=None):
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError):
        return default
    return data


def _write_json(path, data):
    """原子写：先写临时文件再 replace，避免半截文件"""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def _valid_qq(value):
    text = str(value).strip()
    return text.isdigit() and MIN_QQ_LEN <= len(text) <= MAX_QQ_LEN


# ==================== 加载 / 落盘 ====================
def _qq_id(item):
    """QQ 绑定项 -> int，取不到返回 None。

    防御性加固（既有代码）：users/*.json 里 qq 列表若混入非 dict 项或非法 id，
    原写法 `int(item.get("id"))` 会抛 AttributeError/KeyError 而**不在** except 覆盖
    范围内，导致整份身份库加载失败。统一走本取值器后脏数据只被忽略。
    """
    if not isinstance(item, dict):
        return None
    try:
        return int(item.get("id"))
    except (TypeError, ValueError):
        return None


def _index_record(rec):
    uid = int(rec["uid"])
    tg = rec.get("tg") or {}
    if tg.get("raw") is not None:
        try:
            _tg_index[int(tg["raw"])] = uid
        except (TypeError, ValueError):
            pass
    for item in (rec.get("qq") or []):
        qq = _qq_id(item)
        if qq is None:
            continue
        _qq_index[qq] = uid


def _reload():
    """以 users/*.json 为权威重建内存索引；index.json 只用于延续 seq，防止 uid 复用"""
    global _users, _tg_index, _qq_index, _pending, _next_seq
    _users, _tg_index, _qq_index = {}, {}, {}
    try:
        names = sorted(os.listdir(_users_dir()))
    except OSError:
        names = []
    for name in names:
        if not name.endswith(".json"):
            continue
        rec = _read_json(os.path.join(_users_dir(), name))
        if not isinstance(rec, dict):
            continue
        try:
            uid = int(rec.get("uid"))
        except (TypeError, ValueError):
            continue
        rec["uid"] = uid
        _users[uid] = rec
    for rec in _users.values():
        _index_record(rec)
    _next_seq = len(_users) + 1
    idx = _read_json(_index_path())
    if isinstance(idx, dict):
        try:
            _next_seq = max(_next_seq, int(idx.get("seq", 0)) + 1)
        except (TypeError, ValueError):
            pass
    _pending = _read_json(_pending_path(), {})
    if not isinstance(_pending, dict):
        _pending = {}
    _prune_pending()
    _save_index()
    _save_pending()


def _save_index():
    if not _base_dir:
        return
    _write_json(_index_path(), {
        "seq": _next_seq - 1,
        "uid_base": UID_BASE,
        "gid_base": GID_BASE,
        "updated": _stamp(),
        "tg": {str(k): v for k, v in sorted(_tg_index.items())},
        "qq": {str(k): v for k, v in sorted(_qq_index.items())},
    })


def _save_user(rec):
    rec["updated_at"] = _stamp()
    _write_json(_user_path(rec["uid"]), rec)


def _save_pending():
    if not _base_dir:
        return
    _write_json(_pending_path(), _pending)


def _prune_pending():
    now = time.time()
    for token in [t for t, item in _pending.items()
                  if now - float(item.get("created", 0)) > PENDING_TTL]:
        _pending.pop(token, None)


def _new_token():
    for _ in range(50):
        token = "%0*d" % (TOKEN_LEN, random.randint(10 ** (TOKEN_LEN - 1), 10 ** TOKEN_LEN - 1))
        if token not in _pending:
            return token
    raise RuntimeError("无法分配绑定确认码")


def _alloc_uid():
    global _next_seq
    uid = UID_BASE + _next_seq - 1
    _next_seq += 1
    return uid


# ==================== 对外接口 ====================
def configure(accounts_dir, admin_tg_ids=None):
    """由入口调用：注入账号映射目录并加载身份库。admin_tg_ids 为 settings.json 的 admin.tg（裸 id）"""
    global _base_dir, _admin_tg_raw
    with _lock:
        _base_dir = os.path.abspath(accounts_dir)
        os.makedirs(_users_dir(), exist_ok=True)
        _admin_tg_raw = set()
        for value in (admin_tg_ids or []):
            try:
                _admin_tg_raw.add(int(value))
            except (TypeError, ValueError):
                pass
        _reload()
    print("[identity] 身份库已加载：%s（用户 %d，管理员名单 %d）"
          % (_base_dir, len(_users), len(_admin_tg_raw)))
    return _base_dir


def accounts_dir():
    return _base_dir


def resolve(channel_name, user_id):
    """把某通道的用户 id 解析为 uid；未绑定返回 None"""
    try:
        value = int(user_id)
    except (TypeError, ValueError):
        return None
    with _lock:
        if channel_name == "tg":
            offset = _tg_offset()
            raw = value - offset if value >= offset else value
            return _tg_index.get(raw)
        if channel_name == "qq":
            return _qq_index.get(value)
    return None


def ensure_tg(tg_ns_id, username=""):
    """TG 账号首次对话自动开卡（幂等）"""
    with _lock:
        ns = int(tg_ns_id)
        offset = _tg_offset()
        raw = ns - offset if ns >= offset else ns
        uid = _tg_index.get(raw)
        if uid is not None:
            rec = _users.get(uid)
            tg = (rec or {}).get("tg") or {}
            if rec and username and tg.get("username") != username:
                tg["username"] = username
                rec["tg"] = tg
                _save_user(rec)
            return uid
        uid = _alloc_uid()
        rec = {
            "uid": uid,
            "created_at": _stamp(),
            "admin": raw in _admin_tg_raw,
            "tg": {"raw": raw, "ns": ns, "username": username,
                   "source": "auto", "bound_at": _stamp()},
            "qq": [],
            "note": "",
        }
        _users[uid] = rec
        _index_record(rec)
        _save_user(rec)
        _save_index()
        return uid


def force_bind_qq(qq, uid, source="manual"):
    """不做确认直接绑定（迁移/管理员操作用）"""
    with _lock:
        rec = _users.get(int(uid))
        if rec is None:
            return False, "身份 id 不存在"
        qq = int(qq)
        owner = _qq_index.get(qq)
        if owner is not None and owner != int(uid):
            return False, "该 QQ 已绑定到其它身份 %d" % owner
        rec.setdefault("qq", [])
        if not any(_qq_id(x) == qq for x in rec["qq"]):
            rec["qq"].append({"id": qq, "source": source, "bound_at": _stamp()})
        _users[int(uid)] = rec
        _index_record(rec)
        _save_user(rec)
        _save_index()
        return True, "已绑定"


def request_qq_bind(qq, uid):
    """QQ 侧发起绑定 → 返回 (ok, 提示, 确认码)"""
    with _lock:
        if not _valid_qq(qq):
            return False, "QQ 号格式不正确", None
        qq = int(qq)
        rec = _users.get(int(uid))
        if rec is None:
            return False, "身份 id 不存在，请核对 TG 侧 .me 显示的 id", None
        owner = _qq_index.get(qq)
        if owner == int(uid):
            return False, "该 QQ 已经绑定在这个身份上了", None
        if owner is not None:
            return False, "该 QQ 已绑定到其它身份，请先在原身份侧解绑", None
        for token, item in _pending.items():
            if int(item.get("qq", 0)) == qq and int(item.get("uid", 0)) == int(uid):
                # 重发即刷新有效期：QQ 侧每次发起都会重新通知 TG，
                # created 不刷新会让「10 分钟内有效」的文案撒谎
                item["created"] = time.time()
                _save_pending()
                return True, "已有待确认请求", token
        token = _new_token()
        _pending[token] = {"qq": qq, "uid": int(uid), "created": time.time()}
        _save_pending()
        return True, "已发送确认请求", token


def confirm_qq_bind(uid, token):
    """TG 侧确认 → 返回 (ok, 提示, qq)"""
    with _lock:
        token = str(token).strip()
        item = _pending.get(token)
        if not item:
            return False, "确认码无效或已过期", None
        if int(item.get("uid", 0)) != int(uid):
            return False, "该确认码不属于你的身份", None
        if time.time() - float(item.get("created", 0)) > PENDING_TTL:
            _pending.pop(token, None)
            _save_pending()
            return False, "确认码已过期，请在 QQ 侧重新发起绑定", None
        qq = int(item["qq"])
        owner = _qq_index.get(qq)
        if owner is not None and owner != int(uid):
            _pending.pop(token, None)
            _save_pending()
            return False, "该 QQ 已被其它身份绑定", None
        rec = _users[int(uid)]
        rec.setdefault("qq", [])
        if not any(_qq_id(x) == qq for x in rec["qq"]):
            rec["qq"].append({"id": qq, "source": "confirm", "bound_at": _stamp()})
        rec["admin"] = is_admin(uid)
        _users[int(uid)] = rec
        _index_record(rec)
        _pending.pop(token, None)
        _save_user(rec)
        _save_index()
        _save_pending()
        return True, "绑定成功", qq


def unbind_qq(uid, qq):
    with _lock:
        if not _valid_qq(qq):
            return False, "QQ 号格式不正确"
        qq = int(qq)
        rec = _users.get(int(uid))
        if rec is None:
            return False, "身份 id 不存在"
        items = rec.get("qq") or []
        if not any(_qq_id(x) == qq for x in items):
            return False, "该 QQ 未绑定在你的身份上"
        rec["qq"] = [x for x in items if _qq_id(x) != qq]
        _users[int(uid)] = rec
        _qq_index.pop(qq, None)
        _save_user(rec)
        _save_index()
        return True, "已解绑 QQ %d" % qq


def pending_for(uid):
    """该身份下待确认的 QQ 绑定请求"""
    with _lock:
        _prune_pending()
        now = time.time()
        result = []
        for token, item in _pending.items():
            if int(item.get("uid", 0)) != int(uid):
                continue
            remain = max(0, int(PENDING_TTL - (now - float(item.get("created", 0)))))
            result.append({"token": token, "qq": int(item.get("qq", 0)),
                           "remain_seconds": remain})
        return sorted(result, key=lambda x: x["qq"])


def get(uid):
    with _lock:
        rec = _users.get(int(uid))
        return json.loads(json.dumps(rec)) if rec else None


def find_qq(qq):
    try:
        value = int(qq)
    except (TypeError, ValueError):
        return None
    with _lock:
        return _qq_index.get(value)


def list_users():
    with _lock:
        return [get(uid) for uid in sorted(_users)]


def is_admin(uid):
    try:
        value = int(uid)
    except (TypeError, ValueError):
        return False
    with _lock:
        rec = _users.get(value)
        if not rec:
            return False
        if rec.get("admin") is True:
            return True
        tg = rec.get("tg") or {}
        try:
            return int(tg.get("raw")) in _admin_tg_raw
        except (TypeError, ValueError):
            return False


def _push_flags(rec):
    """个人推送开关（默认开）：{"qq": bool, "tg": bool}"""
    push = (rec or {}).get("push") or {}
    return {
        "qq": push.get("qq") is not False,
        "tg": push.get("tg") is not False,
    }


def set_push(uid, platform, enabled):
    """设置某身份的单侧推送开关（.my qq-push / .my tg-push）。返回 (ok, message)。"""
    try:
        value = int(uid)
    except (TypeError, ValueError):
        return False, "身份 id 无效"
    platform = str(platform or "").lower()
    if platform not in ("qq", "tg"):
        return False, "平台只能是 qq 或 tg"
    with _lock:
        rec = _users.get(value)
        if not rec:
            return False, "身份不存在"
        push = rec.get("push") or {}
        push[platform] = bool(enabled)
        rec["push"] = push
        _save_user(rec)
    return True, ("开启" if enabled else "关闭")


def push_targets(uid):
    """统一异步推送的**全部**私聊目标：TG 私聊（若该身份有 TG）+ 每个已绑定 QQ 私聊。

    返回 [("tg", ns), ("qq", raw_qq), ...]，TG 在前、QQ 按绑定顺序、同 uid 内去重。
    身份不存在 / 无任何合法目标时返回 []（不抛异常，由调用方记日志）。
    TG 目标取 rec["tg"]["ns"]；极少数历史记录缺 ns 时按 _tg_offset() 用 raw 回退。
    """
    try:
        value = int(uid)
    except (TypeError, ValueError):
        return []
    with _lock:
        rec = _users.get(value)
        if not rec:
            return []
        targets = []
        tg = rec.get("tg") or {}
        ns = tg.get("ns")
        if ns is None and tg.get("raw") is not None:
            try:
                ns = int(tg["raw"]) + _tg_offset()
            except (TypeError, ValueError):
                ns = None
        if ns is not None:
            try:
                targets.append(("tg", int(ns)))
            except (TypeError, ValueError):
                pass
        seen = set()
        for item in (rec.get("qq") or []):
            if not isinstance(item, dict):
                continue
            raw = item.get("id")
            if raw is None or not _valid_qq(raw):
                continue
            try:
                qq = int(raw)
            except (TypeError, ValueError):
                continue
            if qq in seen:
                continue
            seen.add(qq)
            targets.append(("qq", qq))
        # 个人推送开关（.my qq-push / .my tg-push，默认开）：关闭的一侧从推送目标剔除
        push = rec.get("push") or {}
        if push.get("tg") is False:
            targets = [t for t in targets if t[0] != "tg"]
        if push.get("qq") is False:
            targets = [t for t in targets if t[0] != "qq"]
        return targets


def push_target(uid):
    """兼容接口：统一异步推送的**主目标**（TG 私聊）。返回 (channel_name, target_id) 或 None。

    需要 TG+QQ 镜像推送的调用方请改用 push_targets() / channel.push_to_user()。
    """
    for name, target_id in push_targets(uid):
        if name == "tg":
            return (name, target_id)
    return None


def summary(uid):
    with _lock:
        rec = _users.get(int(uid))
        if not rec:
            return None
        tg = rec.get("tg") or {}
        return {
            "uid": int(rec["uid"]),
            "created_at": rec.get("created_at") or "",
            "admin": is_admin(uid),
            "tg_raw": tg.get("raw"),
            "tg_ns": tg.get("ns"),
            "tg_username": tg.get("username") or "",
            "qq": [q for q in (_qq_id(x) for x in (rec.get("qq") or [])) if q is not None],
            "pending": pending_for(uid),
            "push": _push_flags(rec),
        }


def admin_uids():
    """所有管理员的统一身份 id（以身份库为准：TG/QQ 两侧管理员已合并到同一 uid）。

    供异步告警按 uid 扇出（push_to_user）使用；不再依赖 settings.json 的
    admin.<channel> 原始 id（QQ 侧那份在统一身份改造后恒为空）。
    返回排序后的 uid 列表；身份库为空/无管理员时返回 []。
    """
    result = []
    with _lock:
        for uid in sorted(_users):
            try:
                if is_admin(uid):
                    result.append(int(uid))
            except (TypeError, ValueError):
                continue
    return result


def stats():
    with _lock:
        return {"users": len(_users), "tg": len(_tg_index), "qq": len(_qq_index),
                "pending": len(_pending), "next_seq": _next_seq}