222 lines
8.5 KiB
Python
222 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
||
"""shell_screen_check.py - 验证 /shell 记录里对全屏程序的处理
|
||
|
||
交互式 shell 的原始输出里混着转义序列、重绘空格和全屏程序的整屏画面。
|
||
不清洗的话模型看到的就是几百个空格加一串光标定位；清过头又会把命令和
|
||
状态行一起吃掉。这里用 printf 造出确定的输入来钉住行为：
|
||
|
||
1. 备用屏幕（vim/less/top 走这条路）——画面必须省略，只留首尾标记；
|
||
2. 成片空格（nano/top 靠空格重绘）——必须压掉，但短缩进要原样保留；
|
||
3. 大输出里夹着的备用屏幕——跨 read 也要能认出来。
|
||
|
||
用法：
|
||
shell_screen_check.py <ligent 二进制> <mock base-url>
|
||
"""
|
||
import glob
|
||
import json
|
||
import os
|
||
import pty
|
||
import re
|
||
import select
|
||
import shutil
|
||
import sys
|
||
import time
|
||
|
||
ALT_TEXT = "备用屏幕里的机密内容"
|
||
NORMAL_TEXT = "正常输出"
|
||
|
||
|
||
def drive(binary, base, commands, wait=1.6, term="xterm-256color"):
|
||
"""在 pty 里跑 ligent，进 /shell 执行 commands，导出后返回记录文本。
|
||
|
||
term=None 表示不设 TERM，用来观察退化行为。"""
|
||
saved = os.environ.copy()
|
||
os.environ["LIGENT_API_KEY"] = "sk-test"
|
||
# 没有 TERM 的话 less/vim 会退回非全屏模式，测不到备用屏幕那条路。
|
||
if term is None:
|
||
os.environ.pop("TERM", None)
|
||
else:
|
||
os.environ["TERM"] = term
|
||
|
||
# 每轮都清空归档：归档名是随机的，按文件名排序取"最后一个"会拿到上一轮。
|
||
bindir = os.path.dirname(os.path.abspath(binary))
|
||
shutil.rmtree(os.path.join(bindir, "ligent.archive"), ignore_errors=True)
|
||
|
||
pid, fd = pty.fork()
|
||
if pid == 0:
|
||
os.execv(binary, [binary, "--base-url", base, "--no-confirm", "--auto-save"])
|
||
os._exit(127)
|
||
|
||
# 行编辑器每次回到提示符都会打开括号粘贴 —— 用它当"轮到输入了"的信号。
|
||
# 只按"输出安静"等是不行的：等待动画每 250ms 画一帧，pty 永远不会安静
|
||
# 超过 idle，于是每次都要白等满整个额度。
|
||
PROMPT_READY = b"\x1b[?2004h"
|
||
ready_seen = 0
|
||
tail = b""
|
||
|
||
def note(chunk):
|
||
nonlocal tail, ready_seen
|
||
buf = tail + chunk
|
||
ready_seen += buf.count(PROMPT_READY)
|
||
tail = buf[-(len(PROMPT_READY) - 1):]
|
||
|
||
def drain(seconds, idle=0.4, until_ready=False):
|
||
want = ready_seen + 1 if until_ready else None
|
||
end = time.time() + seconds
|
||
last = time.time()
|
||
while time.time() < end:
|
||
r, _, _ = select.select([fd], [], [], 0.02)
|
||
if r:
|
||
try:
|
||
chunk = os.read(fd, 65536)
|
||
except OSError:
|
||
return
|
||
if not chunk:
|
||
return
|
||
note(chunk)
|
||
last = time.time()
|
||
if want is not None and ready_seen >= want:
|
||
return
|
||
elif time.time() - last > idle:
|
||
return
|
||
|
||
try:
|
||
drain(1.5, until_ready=True) # 等横幅画完、回到提示符
|
||
for line in ["/shell"] + commands + ["exit", "/export", "/exit"]:
|
||
os.write(fd, line.encode("utf-8") + b"\n")
|
||
drain(wait, until_ready=True)
|
||
finally:
|
||
for fn in (lambda: os.close(fd), lambda: os.waitpid(pid, 0)):
|
||
try:
|
||
fn()
|
||
except OSError:
|
||
pass
|
||
os.environ.clear()
|
||
os.environ.update(saved)
|
||
|
||
bindir = os.path.dirname(os.path.abspath(binary))
|
||
files = sorted(glob.glob(os.path.join(bindir, "ligent.archive", "*.json")))
|
||
if not files:
|
||
return None
|
||
with open(files[0], encoding="utf-8") as fh:
|
||
data = json.load(fh)
|
||
users = [m["content"] for m in data["messages"] if m["role"] == "user"]
|
||
return users[-1] if users else ""
|
||
|
||
|
||
def between_markers(rec, start="[全屏程序运行中，画面内容已省略]", end="[全屏程序结束]"):
|
||
"""取出"全屏程序"首尾标记之间的内容。
|
||
|
||
命令回显必然把输入里的字面文本也录进去，所以不能直接断言"机密文本不在
|
||
记录里"——只能断言它没出现在真正被省略的那段区间内。"""
|
||
parts = rec.split(start)
|
||
if len(parts) < 2:
|
||
return None
|
||
return parts[1].split(end)[0]
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 3:
|
||
print("用法: shell_screen_check.py <二进制> <base-url>")
|
||
return 2
|
||
binary, base = sys.argv[1], sys.argv[2]
|
||
bindir = os.path.dirname(os.path.abspath(binary))
|
||
archive = os.path.join(bindir, "ligent.archive")
|
||
|
||
fails = 0
|
||
total = 0
|
||
|
||
def check(name, ok, detail=""):
|
||
nonlocal fails, total
|
||
total += 1
|
||
if ok:
|
||
print(" ✓ %s" % name)
|
||
else:
|
||
fails += 1
|
||
print(" ✗ %s%s" % (name, (" " + detail) if detail else ""))
|
||
|
||
# ---- 1. 备用屏幕：画面省略，只留首尾标记 ----
|
||
print("备用屏幕（vim/less 路径）")
|
||
rec = drive(binary, base, [
|
||
r"printf '\033[?1049h" + ALT_TEXT + r"\033[?1049l'",
|
||
"echo " + NORMAL_TEXT,
|
||
])
|
||
if rec is None:
|
||
check("拿到记录", False, "没有生成归档")
|
||
return 1
|
||
check("拿到记录", True)
|
||
body = between_markers(rec)
|
||
check("首尾标记成对出现", body is not None, "没有找到完整的一对标记")
|
||
check("全屏区间内没有画面内容", body is not None and ALT_TEXT not in body,
|
||
"机密内容漏出来了")
|
||
check("留下开始标记", "全屏程序运行中" in rec)
|
||
check("留下结束标记", "全屏程序结束" in rec)
|
||
check("标记之外的输出还在", NORMAL_TEXT in rec)
|
||
check("记录里没有裸转义符", "\x1b" not in rec)
|
||
|
||
# ---- 2. 成片空格：压掉，但缩进保留 ----
|
||
print("成片空格（nano/top 重绘路径）")
|
||
rec = drive(binary, base, [
|
||
r"printf 'A'; printf '%0.s ' $(seq 1 200); printf 'B\n'",
|
||
r"printf ' 缩进四个空格\n'",
|
||
])
|
||
if rec is None:
|
||
check("拿到记录", False, "没有生成归档")
|
||
return 1
|
||
check("拿到记录", True)
|
||
check("200 个空格被压缩", not re.search(r" {9,}", rec), "还留着成片空白")
|
||
check("压缩后仍保留 A 和 B", "A" in rec and "B" in rec)
|
||
check("短缩进原样保留", " 缩进四个空格" in rec, "缩进被误伤")
|
||
blank = sum(1 for ch in rec if ch in " \n\t") / max(1, len(rec))
|
||
check("空白占比低于 40%", blank < 0.40, "实际 %.0f%%" % (100 * blank))
|
||
|
||
# ---- 3. 大输出里夹着的备用屏幕 ----
|
||
print("大输出（跨 read）里的备用屏幕")
|
||
rec = drive(binary, base, [
|
||
"head -c 8000 /dev/zero | tr '\\0' 'x'; printf '\\033[?1049h"
|
||
+ ALT_TEXT + "\\033[?1049l'",
|
||
])
|
||
if rec is None:
|
||
check("拿到记录", False, "没有生成归档")
|
||
return 1
|
||
check("拿到记录", True)
|
||
body = between_markers(rec)
|
||
check("跨块后仍认得出备用屏幕", body is not None and ALT_TEXT not in body,
|
||
"切断后漏了画面")
|
||
check("跨块标记成对出现", body is not None)
|
||
|
||
# ---- 4. 同一个会话里进两次全屏，标记必须成对 ----
|
||
print("连续两次进入全屏")
|
||
rec = drive(binary, base, [
|
||
r"printf '\033[?1049hAAA\033[?1049l'; printf '\033[?1049hBBB\033[?1049l'",
|
||
])
|
||
if rec is None:
|
||
check("拿到记录", False, "没有生成归档")
|
||
return 1
|
||
check("拿到记录", True)
|
||
check("两次开始标记", rec.count("[全屏程序运行中，画面内容已省略]") == 2,
|
||
"实际 %d 个" % rec.count("[全屏程序运行中，画面内容已省略]"))
|
||
check("两次结束标记", rec.count("[全屏程序结束]") == 2,
|
||
"实际 %d 个" % rec.count("[全屏程序结束]"))
|
||
|
||
# ---- 5. TERM 退化时补一个像样的值 ----
|
||
# TERM=dumb 时 nano 会放弃备用屏幕、改成逐行重绘，整段忽略就失效了。
|
||
# 而 ligent 自己的 REPL 就在用光标定位，说明终端并非真的 dumb。
|
||
for bad in ("dumb", None):
|
||
label = "TERM=%s" % (bad if bad else "(未设置)")
|
||
print("退化的 %s" % label)
|
||
rec = drive(binary, base, ["echo TERM_IS=$TERM"], term=bad)
|
||
if rec is None:
|
||
check("拿到记录", False, "没有生成归档")
|
||
continue
|
||
check("%s 时补上可用 TERM" % label, "TERM_IS=xterm-256color" in rec,
|
||
"实际记录：%s" % rec.strip().replace("\n", " | ")[-80:])
|
||
|
||
shutil.rmtree(archive, ignore_errors=True)
|
||
print("\n%d 项检查，%d 项失败" % (total, fails))
|
||
return 1 if fails else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|