ssh-e2e/storage.py

"""配置存储:JSON 文件 + 平台相关的机密加密。
Windows: DPAPI(绑定当前用户)
Linux/macOS: keyring(可用则用,不可用则明文并在加载时告警)
"""
from __future__ import annotations

import json
import os
import sys
import threading
import uuid
from pathlib import Path

_SENSITIVE = ("password", "key_passphrase")
_PREFIX = "dpapi:"       # Windows DPAPI
_PREFIX2 = "kr:"         # keyring,base64(AES-GCM)


def data_root() -> Path:
    """配置目录:Windows %APPDATA%\\SshTunnelManager / macOS ~/ssh-e2e/Library / Linux ~/ssh-e2e/.config"""
    app = "SshTunnelManager"
    if sys.platform == "win32":
        base = os.environ.get("APPDATA") or str(Path.home() / "AppData" / "Roaming")
        return Path(base) / app
    if sys.platform == "darwin":
        return Path.home() / "ssh-e2e" / "Library"
    base = os.environ.get("XDG_CONFIG_HOME")
    if base and Path(base).is_absolute():
        return Path(base)
    return Path.home() / "ssh-e2e" / ".config"


# ---------- Windows DPAPI ----------

def _dpapi_protect(text: str) -> str:
    import base64
    import ctypes
    from ctypes import wintypes

    class _DataBlob(ctypes.Structure):
        _fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.c_void_p)]

    raw = text.encode("utf-8")
    buf = ctypes.create_string_buffer(raw, len(raw))
    src = _DataBlob(len(raw), ctypes.cast(buf, ctypes.c_void_p))
    dst = _DataBlob()
    ok = ctypes.windll.crypt32.CryptProtectData(
        ctypes.byref(src), None, None, None, None, 0x1, ctypes.byref(dst))
    if not ok:
        raise OSError("DPAPI 加密失败")
    try:
        return _PREFIX + base64.b64encode(
            ctypes.string_at(dst.pbData, dst.cbData)).decode("ascii")
    finally:
        ctypes.windll.kernel32.LocalFree(ctypes.c_void_p(dst.pbData))


def _dpapi_unprotect(text: str) -> str:
    import base64
    import ctypes
    from ctypes import wintypes

    class _DataBlob(ctypes.Structure):
        _fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.c_void_p)]

    raw = base64.b64decode(text[len(_PREFIX):])
    buf = ctypes.create_string_buffer(raw, len(raw))
    src = _DataBlob(len(raw), ctypes.cast(buf, ctypes.c_void_p))
    dst = _DataBlob()
    ok = ctypes.windll.crypt32.CryptUnprotectData(
        ctypes.byref(src), None, None, None, None, 0, ctypes.byref(dst))
    if not ok:
        raise OSError("DPAPI 解密失败:配置可能属于其他 Windows 用户")
    try:
        return ctypes.string_at(dst.pbData, dst.cbData).decode("utf-8")
    finally:
        ctypes.windll.kernel32.LocalFree(ctypes.c_void_p(dst.pbData))


# ---------- keyring(Linux/macOS) ----------

def _keyring():
    try:
        import keyring  # noqa: PLC0415
        if keyring.get_keyring().priority >= 0:  # 有真实后端(非 fail/null)
            return keyring
    except Exception:
        pass
    return None


def _kr_protect(text: str) -> str:
    import base64
    kr = _keyring()
    if kr is None:
        return text  # 无后端:明文(加载时会告警)
    import hashlib
    key = hashlib.sha256((kr.get_password("ssh-tunnel-manager", "salt")
                          or "fixed-salt").encode()).digest()[:16]
    nonce = os.urandom(12)
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    ct = AESGCM(key).encrypt(nonce, text.encode("utf-8"), None)
    return _PREFIX2 + base64.b64encode(nonce + ct).decode("ascii")


def _kr_unprotect(text: str) -> str:
    if not text.startswith(_PREFIX2):
        return text
    import base64
    raw = base64.b64decode(text[len(_PREFIX2):])
    kr = _keyring()
    if kr is None:
        return ""
    import hashlib
    key = hashlib.sha256((kr.get_password("ssh-tunnel-manager", "salt")
                          or "fixed-salt").encode()).digest()[:16]
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    return AESGCM(key).decrypt(raw[:12], raw[12:], None).decode("utf-8")


# ---------- 统一入口 ----------

def protect(text: str) -> str:
    if not text:
        return text
    if sys.platform == "win32":
        return _dpapi_protect(text)
    return _kr_protect(text)


def unprotect(text: str) -> str:
    if not isinstance(text, str) or not text:
        return text
    if text.startswith(_PREFIX):
        if sys.platform != "win32":
            raise OSError("该配置来自 Windows 的 DPAPI 加密,无法在当前平台解密")
        return _dpapi_unprotect(text)
    if text.startswith(_PREFIX2):
        return _kr_unprotect(text)
    return text


class Store:
    """hosts/tunnels 配置的读写。内存中敏感字段为明文,落盘前加密。"""

    def __init__(self, data_dir: Path):
        self.data_dir = data_dir
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.path = data_dir / "config.json"
        self._lock = threading.Lock()
        self.warn_plaintext = False
        self.data = self._load()

    def _load(self) -> dict:
        data = {"hosts": [], "tunnels": []}
        if self.path.exists():
            try:
                loaded = json.loads(self.path.read_text(encoding="utf-8"))
                data["hosts"] = loaded.get("hosts") or []
                data["tunnels"] = loaded.get("tunnels") or []
            except Exception:
                try:
                    self.path.rename(self.path.with_suffix(".json.broken"))
                except OSError:
                    pass
        for h in data["hosts"]:
            for k in _SENSITIVE:
                v = h.get(k)
                if isinstance(v, str) and v:
                    if v.startswith((_PREFIX, _PREFIX2)):
                        try:
                            h[k] = unprotect(v)
                        except Exception:
                            h[k] = ""
                    else:
                        self.warn_plaintext = True  # 明文机密(如无 keyring 的 Linux)
        return data

    def save(self) -> None:
        with self._lock:
            out = json.loads(json.dumps(self.data))  # 深拷贝
            for h in out["hosts"]:
                for k in _SENSITIVE:
                    if h.get(k):
                        h[k] = protect(h[k])
            tmp = self.path.with_suffix(".json.tmp")
            tmp.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
            os.replace(tmp, self.path)

    @staticmethod
    def new_id() -> str:
        return uuid.uuid4().hex[:8]