#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""S3 取证用只读 SQL 探针（可选 --drop-sandbox 收尾清理）。
task OqAv27u3w8DE9nirTwPp2（S3-c develop 自有 commit 收口标注）：本探针随本轮 develop 自有提交纳入版本控制（非引擎代收口），与文档 §7 声明的文件清单一致。

在沙箱库上逐条执行 SELECT 并原样回显结果，供取证链记录"重放前后 count(*)、
主键长度回显、trigger 终态"等事实。凭据唯一事实源 =
<workspace>/projects/pbls/env/test.json 的 db.sandbox（scope=sandbox_only），
口令/账号在输出中一律脱敏为 ******。

用法::

    python3 tests/s3_sql_probe.py "SELECT COUNT(*) AS c FROM pbl_evidence" "SHOW TABLES"
    python3 tests/s3_sql_probe.py --schema pbl_m5a_u7rb --drop-sandbox

安全约束（硬）：非 --drop-sandbox 模式下只允许 SELECT / SHOW / DESC(RIBE) /
information_schema 查询；出现 DML/DDL 关键字直接拒绝并非 0 退出，避免取证脚本
变成改数据的后门。--drop-sandbox 只允许 DROP 掉 env 里声明的 sandbox_schema。
"""
import argparse
import json
import pathlib
import re
import sys

import pymysql

TESTS_DIR = pathlib.Path(__file__).resolve().parent
REPO_ROOT = TESTS_DIR.parent
WORKSPACE_ROOT = REPO_ROOT.parent.parent
ENV_FILE = WORKSPACE_ROOT / "projects" / "pbls" / "env" / "test.json"

FORBIDDEN = re.compile(
    r"\b(INSERT|UPDATE|DELETE|REPLACE|CREATE|ALTER|DROP|TRUNCATE|GRANT|REVOKE|MERGE)\b",
    re.IGNORECASE)
ALLOWED_HEAD = re.compile(r"^\s*(SELECT|SHOW|DESC|DESCRIBE|EXPLAIN)\b", re.IGNORECASE)


def load_sandbox():
    cfg = json.loads(ENV_FILE.read_text(encoding="utf-8"))
    sb = cfg["db"]["sandbox"]
    if sb.get("scope") != "sandbox_only":
        raise SystemExit("refuse non-sandbox scope: %r" % sb.get("scope"))
    return sb


def mask(sb, text):
    out = str(text)
    for secret in (sb.get("password"), sb.get("user")):
        if secret:
            out = out.replace(str(secret), "******")
    return out


def main(argv=None):
    ap = argparse.ArgumentParser(
        description="S3 取证只读 SQL 探针（默认拒写；--drop-sandbox 仅清理沙箱 schema）")
    ap.add_argument("sqls", nargs="*", help="要执行并原样回显的 SELECT/SHOW 语句（可多条）")
    ap.add_argument("--schema", default=None,
                    help="沙箱 schema（缺省取 env/test.json db.sandbox.sandbox_schema）")
    ap.add_argument("--drop-sandbox", action="store_true",
                    help="取证结束后 DROP DATABASE 沙箱 schema（不影响任何其他库）")
    ns = ap.parse_args(argv)

    sb = load_sandbox()
    schema = ns.schema or sb["sandbox_schema"]
    rc = 0

    # 查询模式连到沙箱 schema（否则 SELECT 报 1046 No database selected）；
    # DROP DATABASE 需要不指定库的根连接，故分开建连。
    def _connect(with_db):
        kw = dict(host=sb["host"], port=int(sb["port"]), user=sb["user"],
                  password=sb["password"], charset="utf8mb4",
                  cursorclass=pymysql.cursors.DictCursor, autocommit=True)
        if with_db:
            kw["database"] = schema
        return pymysql.connect(**kw)

    conn = _connect(bool(ns.sqls))
    try:
        for sql in ns.sqls:
            if not ALLOWED_HEAD.match(sql) or FORBIDDEN.search(sql):
                print("REJECT(只读门禁): %s" % sql)
                rc = 1
                continue
            print("SQL> %s" % sql)
            with conn.cursor() as cur:
                cur.execute(sql)
                rows = list(cur.fetchall() or [])
            print("  -> %d 行" % len(rows))
            for r in rows:
                print("  | " + json.dumps({k: mask(sb, v) for k, v in r.items()},
                                          ensure_ascii=False))
        if ns.drop_sandbox:
            conn.close()
            conn = _connect(False)
            print("SQL> DROP DATABASE IF EXISTS `%s`   (仅沙箱 schema，凭据不落盘)" % schema)
            with conn.cursor() as cur:
                cur.execute("DROP DATABASE IF EXISTS `%s`" % schema)
            print("DROP DATABASE %s 完成" % schema)
    except Exception as exc:                          # noqa: BLE001
        print("ERROR: %s: %s" % (type(exc).__name__, mask(sb, exc)))
        rc = 1
    finally:
        conn.close()
    return rc


if __name__ == "__main__":
    sys.exit(main())