"""Linux 微信进程内存密钥提取器。

参考 wechat-cli-main 实现：
- 通过 /proc/<pid>/maps 枚举微信进程可读内存段
- 通过 /proc/<pid>/mem 读取内存
- 匹配 x'<hex>' 模式（64 位 enc_key、96 位 enc_key+salt、更长 hex）
- 用 HMAC-SHA512 校验确认 enc_key
- 返回 salt_hex -> enc_key_hex 映射
"""

from __future__ import annotations

import hashlib
import hmac as hmac_mod
import logging
import os
import re
import struct


_logger = logging.getLogger(__name__)

PAGE_SZ = 4096
KEY_SZ = 32
SALT_SZ = 16
IV_SIZE = 16
HMAC_SIZE = 64
RESERVE_SIZE = 80
MAC_SALT_XOR = 0x3A

# 单次内存读取上限（避免超大映射导致内存爆炸）
_MAX_REGION_SIZE = 500 * 1024 * 1024

_KNOWN_COMMS = {"wechat", "wechatappex", "weixin"}
# crashpad / 渲染辅助子进程：内存中不含主进程的 DB 密钥，扫描会浪费时间
_EXCLUDE_COMM_KEYWORDS = ("crashpad", "crashpad_handler", "renderer", "gpu-process")
_INTERPRETER_PREFIXES = ("python", "bash", "sh", "zsh", "node", "perl", "ruby")
_SKIP_MAPPINGS = {"[vdso]", "[vsyscall]", "[vvar]"}
_SKIP_PATH_PREFIXES = ("/usr/lib/", "/lib/", "/usr/share/")


def _safe_readlink(path: str) -> str:
    try:
        return os.path.realpath(os.readlink(path))
    except OSError:
        return ""


def _is_wechat_process(pid: int) -> bool:
    """检查 pid 是否为微信主进程（排除 crashpad 等辅助子进程）。

    排除逻辑：
    - crashpad_handler / crashpad 子进程：内存中无 DB 密钥，扫描浪费时间
    - 渲染/GPU 辅助进程：同上
    - 解释器进程（python/bash 等）：避免误中
    """
    if pid == os.getpid():
        return False
    try:
        with open(f"/proc/{pid}/comm", encoding="utf-8", errors="replace") as f:
            comm = f.read().strip()
        comm_lower = comm.lower()
        # 排除 crashpad 等辅助子进程
        if any(kw in comm_lower for kw in _EXCLUDE_COMM_KEYWORDS):
            return False
        if comm_lower in _KNOWN_COMMS:
            return True
        exe_path = _safe_readlink(f"/proc/{pid}/exe")
        exe_name = os.path.basename(exe_path)
        if any(exe_name.lower().startswith(p) for p in _INTERPRETER_PREFIXES):
            return False
        return "wechat" in exe_name.lower() or "weixin" in exe_name.lower()
    except (PermissionError, FileNotFoundError, ProcessLookupError):
        return False


def _get_wechat_pids() -> list[tuple[int, int]]:
    """返回疑似微信进程的 (pid, rss_kb) 列表，按内存占用降序。"""
    pids: list[tuple[int, int]] = []
    for pid_str in os.listdir("/proc"):
        if not pid_str.isdigit():
            continue
        pid = int(pid_str)
        try:
            if not _is_wechat_process(pid):
                continue
            with open(f"/proc/{pid}/statm", encoding="utf-8") as f:
                rss_pages = int(f.read().split()[1])
            rss_kb = rss_pages * 4
            pids.append((pid, rss_kb))
        except (PermissionError, FileNotFoundError, ProcessLookupError, OSError, ValueError):
            continue
    pids.sort(key=lambda x: x[1], reverse=True)
    return pids


def _get_readable_regions(pid: int, full_scan: bool = False) -> list[tuple[int, int, str]]:
    """解析 /proc/<pid>/maps，返回可读内存区域 (base, size, name) 列表。

    Args:
        pid: 目标进程 PID
        full_scan: 为 True 时不跳过系统库路径（更慢但更全）；默认按已知路径过滤
    """
    regions: list[tuple[int, int, str]] = []
    maps_path = f"/proc/{pid}/maps"
    with open(maps_path, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            parts = line.split()
            if len(parts) < 2:
                continue
            perms = parts[1]
            if "r" not in perms:
                continue
            mapping_name = parts[5] if len(parts) >= 6 else ""
            if mapping_name in _SKIP_MAPPINGS:
                continue
            if not full_scan:
                mapping_lower = mapping_name.lower()
                if (
                    any(mapping_name.startswith(p) for p in _SKIP_PATH_PREFIXES)
                    and "wcdb" not in mapping_lower
                    and "wechat" not in mapping_lower
                    and "weixin" not in mapping_lower
                ):
                    continue
            start_s, end_s = parts[0].split("-")
            start = int(start_s, 16)
            size = int(end_s, 16) - start
            if 0 < size < _MAX_REGION_SIZE:
                regions.append((start, size, mapping_name))
    return regions


def _collect_db_files(db_dir: str) -> tuple[list[tuple[str, str, int, str, bytes]], dict[str, list[str]]]:
    """遍历 db_dir 收集所有 .db 文件及其 salt。

    Returns:
        db_files: [(rel_path, abs_path, size, salt_hex, page1_bytes), ...]
        salt_to_dbs: {salt_hex: [rel_path, ...]}
    """
    db_files: list[tuple[str, str, int, str, bytes]] = []
    salt_to_dbs: dict[str, list[str]] = {}
    if not os.path.isdir(db_dir):
        return db_files, salt_to_dbs
    for root, _, files in os.walk(db_dir):
        for name in files:
            if not name.endswith(".db") or name.endswith("-wal") or name.endswith("-shm"):
                continue
            path = os.path.join(root, name)
            try:
                size = os.path.getsize(path)
            except OSError:
                continue
            if size < PAGE_SZ:
                continue
            try:
                with open(path, "rb") as f:
                    page1 = f.read(PAGE_SZ)
            except OSError:
                continue
            if page1.startswith(b"SQLite format 3\x00"):
                continue
            rel = os.path.relpath(path, db_dir)
            salt = page1[:SALT_SZ].hex()
            db_files.append((rel, path, size, salt, page1))
            salt_to_dbs.setdefault(salt, []).append(rel)
    return db_files, salt_to_dbs


def _verify_enc_key(enc_key: bytes, page1: bytes) -> bool:
    """通过 HMAC-SHA512 校验 page 1 验证 enc_key 是否正确。"""
    salt = page1[:SALT_SZ]
    mac_salt = bytes(b ^ MAC_SALT_XOR for b in salt)
    mac_key = hashlib.pbkdf2_hmac("sha512", enc_key, mac_salt, 2, dklen=KEY_SZ)
    hmac_data = page1[SALT_SZ: PAGE_SZ - RESERVE_SIZE + IV_SIZE]
    stored_hmac = page1[PAGE_SZ - HMAC_SIZE: PAGE_SZ]
    hm = hmac_mod.new(mac_key, hmac_data, hashlib.sha512)
    hm.update(struct.pack("<I", 1))
    return hm.digest() == stored_hmac


class KeyExtractor:
    """Linux 微信进程内存密钥提取器。"""

    def __init__(self, wechat_pid: int, probe_db_path: str):
        """
        Args:
            wechat_pid: 微信进程 PID（为 0 时自动检测）
            probe_db_path: 用于验证 key 的加密 DB 文件路径，或包含 DB 的目录
        """
        self.wechat_pid = wechat_pid
        self.probe_db_path = probe_db_path
        self._hex_re = re.compile(rb"x'([0-9a-fA-F]{64,192})'")

    # ------------------------------------------------------------------
    # 旧接口兼容：返回 probe DB 对应 salt 的 enc_key
    # ------------------------------------------------------------------
    def extract_key(self) -> str | None:
        """扫描内存提取与 probe_db_path 匹配的 SQLCipher 密钥。

        Returns:
            64 位十六进制 enc_key 字符串，未找到返回 None
        """
        db_dir = self.probe_db_path
        if os.path.isfile(db_dir):
            db_dir = os.path.dirname(db_dir)
        key_map = self.extract_all_keys(db_dir)

        # 优先匹配 probe_db_path 自己的 salt
        probe_salt = None
        if os.path.isfile(self.probe_db_path):
            try:
                with open(self.probe_db_path, "rb") as f:
                    page1 = f.read(PAGE_SZ)
                if not page1.startswith(b"SQLite format 3\x00"):
                    probe_salt = page1[:SALT_SZ].hex()
            except OSError:
                pass

        if probe_salt and probe_salt in key_map:
            return key_map[probe_salt]

        # 否则返回任意一个找到的 key
        if key_map:
            return next(iter(key_map.values()))
        return None

    # ------------------------------------------------------------------
    # 新接口：返回 salt_hex -> enc_key_hex 映射
    # ------------------------------------------------------------------
    def extract_all_keys(
        self,
        db_dir: str | None = None,
        full_scan: bool = False,
    ) -> dict[str, str]:
        """扫描内存，为 db_dir 下每个加密 DB salt 提取 enc_key。

        Args:
            db_dir: 微信数据库目录。默认为 probe_db_path 所在目录。
            full_scan: 是否完整扫描所有可读内存区域（默认优先扫描 [heap]，未找全再扫其它）。

        Returns:
            dict: {salt_hex: enc_key_hex}
        """
        if db_dir is None:
            db_dir = self.probe_db_path
            if os.path.isfile(db_dir):
                db_dir = os.path.dirname(db_dir)

        db_files, salt_to_dbs = _collect_db_files(db_dir)
        if not db_files:
            _logger.warning("在 %s 下未找到加密 DB 文件", db_dir)
            return {}

        _logger.info(
            "找到 %d 个加密 DB，%d 个不同 salt",
            len(db_files),
            len(salt_to_dbs),
        )

        if self.wechat_pid:
            pids = [(self.wechat_pid, 0)]
        else:
            pids = _get_wechat_pids()
            if not pids:
                _logger.warning("未检测到微信进程")
                return {}
            _logger.info(
                "检测到 %d 个微信进程候选: %s",
                len(pids),
                [p for p, _ in pids],
            )

        key_map: dict[str, str] = {}
        remaining = set(salt_to_dbs.keys())

        for pid, _rss_kb in pids:
            if not remaining:
                break
            try:
                regions = _get_readable_regions(pid, full_scan=full_scan)
            except (PermissionError, FileNotFoundError, ProcessLookupError) as e:
                _logger.warning("无法读取 /proc/%s/maps: %s", pid, e)
                continue

            if not full_scan:
                # heap 优先，其次是 stack，再按区域大小降序
                def _heap_first_key(item: tuple[int, int, str]) -> tuple[int, int, str]:
                    name = item[2].lower()
                    if name == "[heap]":
                        return (0, 0, name)
                    if name == "[stack]":
                        return (1, 0, name)
                    return (2, -item[1], name)

                regions.sort(key=_heap_first_key)

            total_size = sum(size for _, size, _ in regions)
            total_mib = total_size / (1024 * 1024)
            _logger.info(
                "PID %s 开始扫描 %d 个区域，合计 %.1f MiB（full_scan=%s）",
                pid,
                len(regions),
                total_mib,
                full_scan,
            )

            try:
                mem = open(f"/proc/{pid}/mem", "rb")
            except (PermissionError, FileNotFoundError, ProcessLookupError) as e:
                _logger.warning("无法打开 /proc/%s/mem: %s", pid, e)
                continue

            try:
                for idx, (base, size, name) in enumerate(regions):
                    if not remaining:
                        break
                    _logger.info(
                        "区域 %d/%d PID=%s [%-16s] 0x%016x size=%.1fMiB 剩余salt=%d",
                        idx + 1,
                        len(regions),
                        pid,
                        name or "?",
                        base,
                        size / (1024 * 1024),
                        len(remaining),
                    )
                    try:
                        self._scan_region(
                            mem, base, size, name,
                            db_files, key_map, remaining, pid,
                        )
                    except (OSError, ValueError) as e:
                        _logger.debug("读取区域失败 PID=%s addr=0x%016x: %s", pid, base, e)
                        continue
                    _logger.info(
                        "完成区域 %d/%d，已找到 %d/%d 个 enc_key",
                        idx + 1,
                        len(regions),
                        len(key_map),
                        len(salt_to_dbs),
                    )
            finally:
                mem.close()

        # 交叉验证：用已找到的 key 验证未匹配的 salt（多 DB 可能共用 enc_key）
        self._cross_verify(db_files, key_map, remaining)

        if key_map:
            _logger.info(
                "成功提取 %d/%d 个 salt 的 enc_key",
                len(key_map),
                len(salt_to_dbs),
            )
        else:
            _logger.warning("扫描完成，未找到任何有效 enc_key")
        return key_map

    def _scan_buffer(
        self,
        data: bytes,
        db_files: list[tuple[str, str, int, str, bytes]],
        key_map: dict[str, str],
        remaining: set[str],
        base_addr: int,
        pid: int,
    ) -> None:
        """扫描一段内存数据，匹配 x'hex' 模式并验证。

        hex 长度分支覆盖：
        - 64：纯 enc_key，尝试匹配所有剩余 salt
        - 96：enc_key(64) + salt(32)，按指定 salt 验证
        - >96 且偶数：取前 64 为 enc_key，末 32 为 salt
        - 66~94 偶数：无明确 salt，按纯 enc_key 尝试所有 salt（覆盖 64~192 全范围）
        """
        for m in self._hex_re.finditer(data):
            hex_str = m.group(1).decode()
            addr = base_addr + m.start()
            hex_len = len(hex_str)

            if hex_len == 64:
                enc_key_hex = hex_str
                self._try_key_all_salts(enc_key_hex, db_files, key_map, remaining, addr, pid)
            elif hex_len == 96:
                enc_key_hex = hex_str[:64]
                salt_hex = hex_str[64:]
                self._try_key(enc_key_hex, salt_hex, db_files, key_map, remaining, addr, pid)
            elif hex_len > 96 and hex_len % 2 == 0:
                enc_key_hex = hex_str[:64]
                salt_hex = hex_str[-32:]
                self._try_key(enc_key_hex, salt_hex, db_files, key_map, remaining, addr, pid)
            elif hex_len > 64 and hex_len < 96 and hex_len % 2 == 0:
                # 66~94 长度：无明确 salt 区，按纯 enc_key 尝试所有 salt
                enc_key_hex = hex_str[:64]
                self._try_key_all_salts(enc_key_hex, db_files, key_map, remaining, addr, pid)

    def _scan_region(
        self,
        mem,
        base: int,
        size: int,
        name: str,
        db_files: list[tuple[str, str, int, str, bytes]],
        key_map: dict[str, str],
        remaining: set[str],
        pid: int,
    ) -> None:
        """流式读取并扫描一个内存区域，按 MiB 输出进度。

        采用 8 MiB 分块 + 256 字节重叠，避免超大区域一次性读入内存，同时保证
        x'hex' 模式跨块边界时也能被匹配。
        """
        chunk_size = 8 * 1024 * 1024
        overlap = 256
        tail = b""
        position = 0
        last_log_mib = -1
        region_mib = size / (1024 * 1024)

        while position < size:
            to_read = min(chunk_size, size - position)
            try:
                mem.seek(base + position)
                chunk = mem.read(to_read)
            except (OSError, ValueError):
                break
            if not chunk:
                break

            data = tail + chunk
            buffer_base = base + position - len(tail)
            self._scan_buffer(data, db_files, key_map, remaining, buffer_base, pid)

            position += len(chunk)
            scanned_mib = position / (1024 * 1024)
            current_mib = int(scanned_mib)
            if current_mib > last_log_mib:
                last_log_mib = current_mib
                _logger.info(
                    "PID %s [%-16s] 已扫描 %.1f/%.1f MiB，已找到 %d 个 enc_key",
                    pid,
                    name or "?",
                    scanned_mib,
                    region_mib,
                    len(key_map),
                )

            tail = data[-overlap:] if len(data) >= overlap else data

    def _try_key(
        self,
        enc_key_hex: str,
        salt_hex: str,
        db_files: list[tuple[str, str, int, str, bytes]],
        key_map: dict[str, str],
        remaining: set[str],
        addr: int,
        pid: int,
    ) -> None:
        """尝试用一个 enc_key 匹配指定 salt。"""
        if salt_hex not in remaining:
            return
        try:
            enc_key = bytes.fromhex(enc_key_hex)
        except ValueError:
            return
        for _rel, _path, _sz, s, page1 in db_files:
            if s != salt_hex:
                continue
            if _verify_enc_key(enc_key, page1):
                key_map[salt_hex] = enc_key_hex
                remaining.discard(salt_hex)
                _logger.info(
                    "找到 enc_key: salt=%s PID=%s addr=0x%016x",
                    salt_hex,
                    pid,
                    addr,
                )
            break

    def _try_key_all_salts(
        self,
        enc_key_hex: str,
        db_files: list[tuple[str, str, int, str, bytes]],
        key_map: dict[str, str],
        remaining: set[str],
        addr: int,
        pid: int,
    ) -> None:
        """尝试用一个 64-hex enc_key 匹配所有剩余 salt。"""
        if not remaining:
            return
        try:
            enc_key = bytes.fromhex(enc_key_hex)
        except ValueError:
            return
        for _rel, _path, _sz, salt_hex, page1 in db_files:
            if salt_hex not in remaining:
                continue
            if _verify_enc_key(enc_key, page1):
                key_map[salt_hex] = enc_key_hex
                remaining.discard(salt_hex)
                _logger.info(
                    "找到 enc_key: salt=%s PID=%s addr=0x%016x",
                    salt_hex,
                    pid,
                    addr,
                )

    def _cross_verify(
        self,
        db_files: list[tuple[str, str, int, str, bytes]],
        key_map: dict[str, str],
        remaining: set[str],
    ) -> None:
        """用已找到的 key 交叉验证剩余 salt。"""
        if not remaining or not key_map:
            return
        for _rel, _path, _sz, salt_hex, page1 in db_files:
            if salt_hex not in remaining:
                continue
            for known_salt, known_key_hex in key_map.items():
                enc_key = bytes.fromhex(known_key_hex)
                if _verify_enc_key(enc_key, page1):
                    key_map[salt_hex] = known_key_hex
                    remaining.discard(salt_hex)
                    _logger.info(
                        "交叉验证成功: salt=%s 复用 salt=%s 的 enc_key",
                        salt_hex,
                        known_salt,
                    )
                    break