#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""bash-output-guard —— `PreToolUse(Bash)`：在**命令执行前**拦下"会灌爆上下文"的读命令。

为什么需要（2026-09-15 实测）
────────────────────────────
对话历史是 **append-only**：一次工具调用在转录里落两条记录 —— `function_call`（命令）
+ `function_call_result`（**输出正文**）；**输出一旦生成就永久留在 messages 里，且每轮全量重发**。
实测某会话 **553 轮 × 平均 42 万 token = 2.33 亿 input**，而 output 仅占 **0.35%**。
⇒ 想真正"把大输出从上下文里去掉"，**唯一的时机就是它被生成之前**（模型无权删自己的历史）。

边界（重要）
──────────
* **只拦"读"类高风险命令**，且**必须给出等价限流写法**（教而非堵）。
* **判不准就放行**（fail-open）—— 本机不是沙箱，但这类命令本身不破坏数据，误拦的代价只是麻烦。
* 急停：env `DSH_OUTPUT_GUARD_OFF=1`，或新建 `<工作区>/.workbuddy/bash-guard.disabled`。
* 命中才写一行日志 `<工作区>/.workbuddy/bash-guard.log`（用于调误报，上限 300 行）。

安装（`settings.json` 的 hooks 段 · **新增一条**）
────────────────────────────────────────────────────────────────────
  "PreToolUse": [ …,
    { "matcher": "Bash|Read",
      "hooks": [{ "type": "command",
                  "command": "\"<python>\" \"<此脚本>\"", "timeout": 10 }] } ]
⚠️ matcher 必须含 **`Read`**：本脚本有两条规则 —— `cat` 大文件（>200 KB）走 **Bash**，
   `Read` 大文件（>400 KB，见 `READ_BIG`）走 **Read 工具**；只挂 `Bash` ⇒ 后者整条规则静默失效。
⚠️ 旧注记"需完全重启才生效"**已不成立**（2026-09-30 实测：改完配置后下一次 Bash 调用即被本钩子
   记录 `entry` 行，未重启宿主）。⇒ 但改完仍应**喂一次模拟载荷端到端自证**。
⛔ **别给本脚本加 `-E`**：`-E` 会屏蔽 `PYTHONUTF8`/`PYTHONIOENCODING` ⇒ stdin 回退 cp936 ⇒
   含中文的 payload 解析失败且**静默 fail-open**（同目录 `stop-dialog-guard.py` 已因此"白排查一天"）。
"""
# [session-mechanism] roots.env 外置（由 install.py 生成；缺失则回落到按位置推导）
def _sm_load_roots():
    import os as _os
    _here = _os.path.dirname(_os.path.abspath(__file__))
    for _up in range(4):
        _p = _os.path.join(_here, *([".."] * _up), "roots.env")
        _p = _os.path.normpath(_p)
        if _os.path.isfile(_p):
            try:
                with open(_p, encoding="utf-8") as _f:
                    for _ln in _f:
                        _ln = _ln.strip()
                        if _ln and not _ln.startswith("#") and "=" in _ln:
                            _k, _v = _ln.split("=", 1)
                            _os.environ.setdefault(_k.strip(), _v.strip())
            except Exception:
                pass
            return
_sm_load_roots()

import io
import json
import os
import re
import sys
import time

# 🔴 2026-10-02：环境定位/体检收敛到 `_env`（同目录优先）——⛔ 不再用 `~` 直拼
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    import _env
except Exception:
    sys.stderr.write('[env] 无法导入 _env.py ⇒ 环境定位不可信\n')
    raise

LOG_REL = os.path.join('.workbuddy', 'bash-guard.log')
# 🔴 2026-09-24：本钩子**没有作用域门禁**，一向全机生效。
#    原 `SCOPE = 'ai1net-dsh-server'` 只喂给日志字段，却会被误读成"作用域外 ⇒ 没生效"（实测误判过一次）⇒ 删除。
# 兜底工作区：脚本位于 <工作区>/dsh-server-docs/07-scripts/ ⇒ 上溯三级
WS_FALLBACK = os.environ.get("DSH_WS_ROOT") or os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# (是否 core, 正则, 名称, 更省的写法)
# ⚠️ 设计原则：**不能全拦** —— 误拦挡住正事 ⇒ 我会来回试/绕路 ⇒ 反而更贵；漏拦只是多花点 token。
#   core=True  = 几乎必然巨大、误拦率≈0 ⇒ **默认就拦**
#   core=False = 经常确实需要全量 ⇒ **仅 hard 模式拦**（默认放行）
RULES = [
    # cat：**按目标文件大小判定**（小文件不拦）—— 这是防日常误伤的关键，见 cat_bigfile()
    (True, '__CAT_BIGFILE__',
     '`cat` 大文件', '改用 `head -30 文件` / `sed -n \'1,30p\' 文件` / `wc -l 文件`'),
    # ⚠️ 2026-09-15 **回放今日 12 个会话的 4049 条真实命令**：本规则命中 62 次，抽样 **29/29 全是窄范围**
    #    （单文件 / 具体目录），**0 条从根**；连 `grep -c` 这种本来只有几行的也被误伤。
    #    旧正则 `\bgrep\b[^|;]*-[a-zA-Z]*[rR]` 还有 bug：`[^|;]*` 会匹配到**路径里的 `-server`**
    #    （`-s`+`erve`+`r`）⇒ 凡路径含 `ai1net-dsh-server` 就必命中。
    #    ⇒ **误伤 >> 收益** ⇒ 降为 loose（仅 `hard` 模式拦），并把正则收紧到「选项紧跟 grep」。
    (False, r'\bgrep\s+-[a-zA-Z]*[rR]\b',
     '递归 grep', '改用 `grep -rn … | head -30`，或 `grep -rc …`（只要计数）'),
    # ⚠️ 只拦 `-R`（递归全树）；**不拦 `ls -la`**（含 l 会误伤日常）；选项必须紧跟 ls（同样防路径误匹配）
    (True, r'\bls\s+-[a-zA-Z]*R\b',
     '`ls -R` 递归全树', '改用 `ls -la 目录 | head -20`，或 `ls 目录 | wc -l`'),
    (True, r'\bfind\s+(/[A-Za-z]|[A-Za-z]:)',
     '`find` 从盘符/根起全树扫', '改用 `find 具体目录 -maxdepth 3 … | head -20`'),
    (True, r'\bjournalctl\b(?!.*(-n\s*\d|head|--since))',
     '`journalctl` 无行数/时间限制', '改用 `journalctl -u X -n 50 --no-pager`（`--since` 也算限流，放行）'),
    (True, r'\bdmesg\b(?!.*head)',
     '`dmesg` 无行数限制', '改用 `dmesg | tail -30`'),
    (False, r'(^|[|;&]\s*)rg\b(?!.*\|)',
     '递归 rg 无管道限流', '改用 `rg … | head -30`，或 `rg -c …`'),
    (False, r'\bgit\s+(log|diff)\b(?!.*(-n\s*\d|--max-count|head))',
     '`git log/diff` 无行数限制', '改用 `git log --oneline -5` / `git diff --stat`'),
    (False, r'\b(du|tree)\b[^|;]*\s(/|[A-Za-z]:)',
     '`du/tree` 从根起', '改用 `du -sh 具体目录` / `tree -L 2 目录 | head -40`'),
]
# 命令里已有限流 ⇒ 放行（只按"最外层"有就好了）
SAFE = re.compile(r'\|\s*(head|tail|wc|grep\s+-c|cut|sed\s+-n|awk|uniq|sort\s+-u)\b')


def _read_stdin():
    """⚠️ 必须走 buffer 显式 UTF-8（`-E` 下 sys.stdin 是 cp936）。"""
    try:
        return sys.stdin.buffer.read().decode('utf-8', 'replace')
    except Exception:
        try:
            return sys.stdin.read()
        except Exception:
            return ''


def _emit(obj):
    data = json.dumps(obj, ensure_ascii=False).encode('utf-8')
    try:
        sys.stdout.buffer.write(data)
        sys.stdout.buffer.flush()
    except Exception:
        try:
            sys.stdout.write(data.decode('utf-8', 'replace'))
            sys.stdout.flush()
        except Exception:
            pass


def log(root, detail):
    try:
        p = os.path.join(root, LOG_REL)
        os.makedirs(os.path.dirname(p), exist_ok=True)
        with io.open(p, 'a', encoding='utf-8', newline='\n') as f:
            f.write('%s\t%s\n' % (time.strftime('%Y-%m-%d %H:%M:%S'), detail))
        lines = io.open(p, encoding='utf-8').read().split('\n')
        if len(lines) > 300:
            io.open(p, 'w', encoding='utf-8', newline='\n').write('\n'.join(lines[-150:]))
    except Exception:
        pass


CAT_BIG = 200 * 1024          # `cat` 目标文件 > 200 KB 才值得拦（小文件放行 ⇒ 防日常误伤）
READ_BIG = 400 * 1024         # `Read` 目标文件 > 400 KB ⇒ 至少几万 token ⇒ 拦（2026-09-15 实测补）
#   ↑ 为什么补 Read：本钩子最初只管 Bash，但实测 `Read` 是**第二大输出源** ——
#     某会话 44 次 / 累计 238 K 字符、单条最大 33 K 字符（≈13 K token）；另一会话 136 次 / 217 K 字符，
#     且**完全没有被任何机制覆盖**。⇒ "想把大输出从上下文里去掉"，就必须覆盖所有会灌大输出的工具。


def read_too_big(payload):
    """`Read` 一个大文件 ⇒ True。取不到大小（文件不存在/相对路径）⇒ False（放行）。"""
    ti = payload.get('tool_input') or {}
    p = ti.get('file_path') or ti.get('path') or ''
    if not isinstance(p, str) or not p:
        return False
    try:
        return os.path.getsize(p) > READ_BIG
    except OSError:
        return False


def cat_bigfile(cmd):
    """`cat <文件>` 且目标文件 > CAT_BIG ⇒ True。取不到大小（相对路径/不存在/管道源）⇒ False（放行）。"""
    m = re.search(r'(?:^|[|;&]\s*)cat\s+(?:-[A-Za-z]+\s+)*([^\s|>;&]+)', cmd)
    if not m:
        return False
    p = m.group(1).strip('"\'')
    if p in ('-', '/dev/null', '/dev/stdin'):
        return False
    try:
        return os.path.getsize(p) > CAT_BIG
    except OSError:
        return False


def segments(cmd):
    """按 shell 分隔符切段（`;` `&&` `||` `|` 换行）—— 只对**每段的开头**做匹配。

    ⚠️ 为什么不匹配整条命令文本（2026-09-15 实测教训）：那样**引号里的字符串也会被拦** ——
    例如 `printf '...ls -laR...'`、测试脚本、把命令写进文档，全都会误拦（阻塞面过大）。
    按段匹配既保住真拦（`cd x && ls -laR` 的第二段以 `ls` 开头 ⇒ 照样拦），又不误伤"只是提到"。
    """
    return [s.strip() for s in re.split(r'(?:&&|\|\||;|\n|\|)', cmd) if s.strip()]


# ── 批量活检测（2026-09-15 用户选 A）─────────────────────────────
#   要打断的是：**"一句话需求 → 上百次工具调用"** —— 每次调用都把「命令 + 结果」永久留在上下文里。
#   实测（会话 7057685c）：逐处改注释 ⇒ Edit 142 次、痕迹 13.4 万字符；若改用一个脚本 ⇒ 约 8 千字符（省 94%）。
#   实现：`PreToolUse` 里给 Edit/Write/Bash 记账；到档位且**同质** ⇒ deny 一刀并给出替代做法（我改路子）。
BATCH_N = 20
BATCH_TOOLS = ('Edit', 'Write', 'Bash')


def _batch_state(root):
    p = os.path.join(root, '.workbuddy', 'batch-guard.json')
    try:
        d = json.load(io.open(p, encoding='utf-8'))
    except Exception:
        d = {}
    return p, d


def _homogeneous(tool, st):
    """同质性：Edit ⇒ 最近 20 次里同一文件出现 ≥3 次（= 同一文件反复改）；Bash ⇒ 命令都很短。"""
    if tool == 'Edit':
        paths = st.get('paths') or []
        if not paths:
            return False
        c = {}
        for x in paths[-20:]:
            c[x] = c.get(x, 0) + 1
        return max(c.values()) >= 3
    if tool == 'Bash':
        lens = st.get('lens') or []
        return bool(lens) and (sum(lens[-20:]) / max(len(lens[-20:]), 1)) < 300
    return True


def batch_check(root, payload, log_fn):
    """到档位 + 同质 ⇒ 返回提示串（并记账，每档只提示一次）；否则 ''。"""
    sid = str(payload.get('session_id') or '')
    tool = payload.get('tool_name') or ''
    if not sid or tool not in BATCH_TOOLS:
        return ''
    p, d = _batch_state(root)
    st = d.setdefault(sid, {})
    st[tool] = int(st.get(tool, 0)) + 1
    n = st[tool]
    ti = payload.get('tool_input') or {}
    if tool == 'Edit':
        fp = str(ti.get('file_path') or '')
        if fp:
            st['paths'] = (st.get('paths') or [])[-40:] + [fp]
    elif tool == 'Bash':
        st['lens'] = (st.get('lens') or [])[-40:] + [len(str(ti.get('command') or ''))]
    fired = st.setdefault('fired', [])
    hit = ''
    if n in (BATCH_N, BATCH_N * 2, BATCH_N * 4) and ('%s@%d' % (tool, n)) not in fired and _homogeneous(tool, st):
        fired.append('%s@%d' % (tool, n))
        files = len(set(st.get('paths') or []))
        hit = (
            '🛠【批量活提示】本会话 `%s` 已调用 **%d 次**%s ⇒ 这更像**批量活**。\n'
            '✅ 改用**一个脚本一次做完**，只回「改了 N 处 / 失败 M 处」—— **别逐处调用**：'
            '逐次调用会把 N 份「命令 + 结果」**永久留在会话上下文里、之后每轮全量重发**。\n'
            'ℹ️ 实测参照：某会话逐处改 142 次 ⇒ 痕迹 **13.4 万字符**；改用一个脚本 ⇒ 约 **8 千字符（省 94%%）**。'
            '改完**统一验一次**即可，不必改一处验一处。'
            % (tool, n, ('（改到的文件只有 %d 个）' % files) if (tool == 'Edit' and files) else ''))
        log_fn(str(root), 'DENY｜批量活提示 %s@%d' % (tool, n))
    try:
        if len(d) > 30:
            for k in list(d)[:-30]:
                d.pop(k, None)
        io.open(p, 'w', encoding='utf-8', newline='\n').write(json.dumps(d, ensure_ascii=False))
    except Exception:
        pass
    return hit


def reason_for(cmd, hard=False):
    segs = segments(cmd)
    for core, pat, why, fix in RULES:
        if not core and not hard:      # loose 条只在 hard 模式生效
            continue
        for s in segs:
            if pat == '__CAT_BIGFILE__':
                if cat_bigfile(s):
                    return why, fix
                continue
            if re.match(pat, s):       # ← match（段首），不是 search（全串）
                return why, fix
    return None, None


def main():
    raw = _read_stdin()
    payload = None
    if raw.strip():
        try:
            payload = json.loads(raw)
        except ValueError:
            payload = None
    root = (os.environ.get('CODEBUDDY_PROJECT_DIR') or os.environ.get('DSH_WORKSPACE')
            or (isinstance(payload, dict) and payload.get('cwd')) or WS_FALLBACK)
    tp = str((payload or {}).get('transcript_path') or '') if isinstance(payload, dict) else ''
    if isinstance(payload, dict):
        # 入口即留痕（只记 event，低频；便于判"有没有被调用"）
        log(str(root), 'entry｜event=%s｜in_scope=all(无门禁)' % (payload.get('hook_event_name') or '(parse-fail)'))
    if not isinstance(payload, dict):
        return
    if (payload.get('hook_event_name') or '') != 'PreToolUse':
        return
    tool = payload.get('tool_name') or ''
    if tool not in ('Bash', 'Read', 'Edit', 'Write'):   # Bash/Read=拦大输出；Edit/Write=批量活记账
        return
    if os.environ.get('DSH_OUTPUT_GUARD_OFF'):
        return
    try:
        if os.path.exists(os.path.join(root, '.workbuddy', 'bash-guard.disabled')):
            return
    except Exception:
        pass
    if tool in ('Edit', 'Write'):     # 批量活检测：deny 一刀，换掉"上百次逐处调用"
        chk = batch_check(str(root), payload, log)
        if chk:
            _emit({'hookSpecificOutput': {'hookEventName': 'PreToolUse',
                                          'permissionDecision': 'deny',
                                          'permissionDecisionReason': chk}})
        return
    if tool == 'Read':
        if read_too_big(payload):
            fp = str(((payload.get('tool_input') or {}).get('file_path')) or '')
            log(str(root), 'DENY｜Read 大文件｜%s' % fp[:110])
            _emit({'hookSpecificOutput': {
                'hookEventName': 'PreToolUse', 'permissionDecision': 'deny',
                'permissionDecisionReason': (
                    '💰 拦下：**Read 大文件**（`%s`）—— 它的**全文会写进会话历史、之后每轮全量重发**。\n'
                    '✅ 分段读：`Read` 加 `offset` / `limit`（例 `limit: 120`）只取你要的那段；'
                    '或先 `grep -n "关键词" 文件 | head -20` 定位行号，再读那几行。\n'
                    'ℹ️ 确实要通读 ⇒ 用 Bash 分批 `sed -n \'1,200p\' 文件`；急停 env `DSH_OUTPUT_GUARD_OFF=1` '
                    '或 `.workbuddy/bash-guard.disabled`。' % fp[:80])}})
        return
    cmd = ((payload.get('tool_input') or {}).get('command')) or ''
    if not cmd or SAFE.search(cmd):
        return
    try:
        md = io.open(os.path.join(root, '.workbuddy', 'bash-guard-mode'),
                     encoding='utf-8').read().lower()
    except OSError:
        md = ''
    if 'off' in md:
        return
    why, fix = reason_for(cmd, hard=('hard' in md))
    if not why:
        return
    log(str(root), 'DENY｜%s｜%s' % (why, cmd.replace('\n', ' ')[:110]))
    _emit({'hookSpecificOutput': {
        'hookEventName': 'PreToolUse',
        'permissionDecision': 'deny',
        'permissionDecisionReason': (
            '💰 拦下：**%s** —— 这类命令的输出会**永久留在会话上下文里、每轮全量重发**'
            '（实测某会话 553 轮 × 平均 42 万 token = 2.33 亿 input，'
            'output 仅占 0.35%%）。\n'
            '✅ 换成限流写法再发：%s\n'
            'ℹ️ 若确实需要全量：**先落盘再只读关键行**（`… > /tmp/x.txt 2>&1` 然后 `sed -n \'1,40p\' /tmp/x.txt`）；'
            '急停用 env `DSH_OUTPUT_GUARD_OFF=1` 或新建 `<工作区>/.workbuddy/bash-guard.disabled`。'
            % (why, fix))}})


if __name__ == '__main__':
    try:
        main()
    except Exception:
        pass          # fail-open：本钩子只为省积分，绝不因自身异常阻断正常工作
    sys.exit(0)