diff --git a/cloud/tests/test_worker_sweep_admin.py b/cloud/tests/test_worker_sweep_admin.py new file mode 100644 index 0000000..d684acd --- /dev/null +++ b/cloud/tests/test_worker_sweep_admin.py @@ -0,0 +1,201 @@ +"""管理员 CLI 清扫接线单测（A5）：鉴权在先，杀在后。 + +只用标准库 + 本地 SQLite 临时库，不联网、不杀任何真实存活 +进程、不碰任何容器、不新增 HTTP 路由：可杀档的“真杀”一律 +用注入的 kill_fn 测试替身（记录调用，不发真实信号）；默认 +审计路径（统一审计口）真实走库，可查 audit_log 验证杀前 / +杀后审计；dry-run 一律不调 reap、不杀。 + +覆盖： +- 坏 key：PermissionError，kill_fn 绝不被调，无杀审计。 +- 已认证但非全局管理员：PermissionError，kill_fn 绝不被调， + 记 admin_worker_sweep_denied。 +- 未配置 CODEX_GLOBAL_ADMINS 时即使是合法用户也 fail-closed。 +- 授权全链路：判定（kill 档）→ 执行时再验 → 杀前审计 → + 杀 → 复核 → 杀后审计 → 汇总审计，outcome 为 killed。 +- 授权 dry-run：只列判定，不杀、不写杀审计。 +- CLI main：授权 exit 0 打 JSON 汇总；被拒 exit 1。 +""" + +import io +import json +import os +import sqlite3 +import sys +import tempfile +import time +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest import mock + +try: + from cloud.auth import AuthStore + from cloud.worker_sweep import main as sweep_main + from cloud.worker_sweep import sweep_stale_workers_admin + from cloud.workers import register_worker +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.auth import AuthStore + from cloud.worker_sweep import main as sweep_main + from cloud.worker_sweep import sweep_stale_workers_admin + from cloud.workers import register_worker + + +def _set_last_seen(db_path, worker_id, last_seen): + conn = sqlite3.connect(db_path) + try: + conn.execute("UPDATE workers SET last_seen = ? WHERE worker_id = ?", + (last_seen, worker_id)) + conn.commit() + finally: + conn.close() + + +def _audit_actions(db_path): + conn = sqlite3.connect(db_path) + try: + return [row[0] for row in conn.execute( + "SELECT action FROM audit_log ORDER BY id").fetchall()] + finally: + conn.close() + + +def _audit_details(db_path, action): + conn = sqlite3.connect(db_path) + try: + return [row[0] for row in conn.execute( + "SELECT detail FROM audit_log WHERE action = ? ORDER BY id", + (action,)).fetchall()] + finally: + conn.close() + + +class WorkerSweepAdminTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + self.now = time.time() + register_worker(self.db_path, "w-kill", host="h-local", pid=9001) + register_worker(self.db_path, "w-foreign", host="h-other", + pid=9003) + register_worker(self.db_path, "w-alive", host="h-local", pid=9004) + _set_last_seen(self.db_path, "w-kill", self.now - 200) + _set_last_seen(self.db_path, "w-foreign", self.now - 200) + _set_last_seen(self.db_path, "w-alive", self.now - 200) + store = AuthStore(self.db_path) + try: + self.admin_id, self.admin_key = store.create_user("admin") + self.user_id, self.user_key = store.create_user("plain") + finally: + store.conn.close() + self.dead = {9001} + self.killed = [] + + def tearDown(self): + self.tmp.cleanup() + + def _checker(self, pid): + return pid in self.dead + + def _kill(self, pid): + self.killed.append(pid) + + def _admin_env(self, admin_id): + return mock.patch.dict(os.environ, + {"CODEX_GLOBAL_ADMINS": admin_id}) + + def test_bad_key_refused_before_any_kill(self): + with self._admin_env(self.admin_id): + with self.assertRaises(PermissionError): + sweep_stale_workers_admin( + self.db_path, "not-a-key", local_host="h-local", + pid_is_dead_fn=self._checker, kill_fn=self._kill, + now=self.now) + self.assertEqual(self.killed, []) + self.assertNotIn("worker_pid_kill", _audit_actions(self.db_path)) + + def test_authenticated_non_admin_refused_and_denial_audited(self): + with self._admin_env(self.admin_id): + with self.assertRaises(PermissionError): + sweep_stale_workers_admin( + self.db_path, self.user_key, local_host="h-local", + pid_is_dead_fn=self._checker, kill_fn=self._kill, + now=self.now) + self.assertEqual(self.killed, []) + actions = _audit_actions(self.db_path) + self.assertIn("admin_worker_sweep_denied", actions) + self.assertNotIn("worker_pid_kill", actions) + + def test_unconfigured_admin_list_fail_closed(self): + with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": ""}): + with self.assertRaises(PermissionError): + sweep_stale_workers_admin( + self.db_path, self.admin_key, local_host="h-local", + pid_is_dead_fn=self._checker, kill_fn=self._kill, + now=self.now) + self.assertEqual(self.killed, []) + + def test_authorized_full_chain(self): + with self._admin_env(self.admin_id): + summary = sweep_stale_workers_admin( + self.db_path, self.admin_key, local_host="h-local", + pid_is_dead_fn=self._checker, kill_fn=self._kill, + now=self.now) + self.assertEqual(self.killed, [9001]) + self.assertEqual(summary["killed"], ["w-kill"]) + self.assertEqual(sorted(summary["listed"]), + ["w-alive", "w-foreign"]) + self.assertTrue(summary["audit_summary"]) + by_id = {row["worker_id"]: row for row in summary["results"]} + self.assertEqual(by_id["w-kill"]["outcome"], "killed") + self.assertTrue(by_id["w-kill"]["audit_before"]) + self.assertTrue(by_id["w-kill"]["audit_after"]) + self.assertTrue(by_id["w-kill"]["verified_dead"]) + self.assertEqual(by_id["w-kill"]["action"], "kill") + actions = _audit_actions(self.db_path) + self.assertIn("worker_pid_kill", actions) + self.assertIn("worker_pid_kill_verified", actions) + self.assertIn("admin_worker_sweep", actions) + kill_details = _audit_details(self.db_path, "worker_pid_kill") + self.assertTrue(any("w-kill" in detail and "9001" in detail + for detail in kill_details)) + + def test_authorized_dry_run_lists_without_killing(self): + with self._admin_env(self.admin_id): + summary = sweep_stale_workers_admin( + self.db_path, self.admin_key, local_host="h-local", + pid_is_dead_fn=self._checker, kill_fn=self._kill, + dry_run=True, now=self.now) + self.assertTrue(summary["dry_run"]) + self.assertEqual(self.killed, []) + self.assertEqual(summary["killed"], []) + by_id = {row["worker_id"]: row for row in summary["results"]} + self.assertEqual(by_id["w-kill"]["action"], "kill") + self.assertNotIn("worker_pid_kill", _audit_actions(self.db_path)) + + def test_cli_main_authorized_and_refused(self): + # 授权路径走 --dry-run：只断言 CLI 鉴权打通、JSON 汇总 + # 结构正确；真杀全链路由上面的注入测试覆盖，不让 CLI + # 测试对假 pid 走真实默认杀路径。 + with self._admin_env(self.admin_id): + buf = io.StringIO() + with redirect_stdout(buf): + code = sweep_main(["--db", self.db_path, + "--api-key", self.admin_key, + "--host", "h-local", "--dry-run"]) + self.assertEqual(code, 0) + payload = json.loads(buf.getvalue()) + self.assertTrue(payload["dry_run"]) + self.assertEqual(payload["killed"], []) + self.assertEqual(sorted(payload["listed"]), + ["w-alive", "w-foreign", "w-kill"]) + with self._admin_env(self.admin_id): + with self.assertRaises(SystemExit) as caught: + sweep_main(["--db", self.db_path, + "--api-key", "not-a-key"]) + self.assertEqual(caught.exception.code, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker_sweep.py b/cloud/worker_sweep.py new file mode 100644 index 0000000..ea6d1f6 --- /dev/null +++ b/cloud/worker_sweep.py @@ -0,0 +1,174 @@ +"""管理员 CLI：失联 worker 清扫（workers kill 接线，A5）。 + +用法：python3 -m cloud.worker_sweep --db <控制面库路径> [--dry-run] + [--stale-seconds N] [--host <本机主机名>] + API key 走环境变量 CLOUD_WORKER_SWEEP_API_KEY，或 --api-key + 显式给（命令行会进进程表，生产优先用环境变量）。 + +这是 reap_stale_workers 的第一个生产调用方（此前除 workers.py +本体与测试外全仓库零生产调用方）： + +- 鉴权：先经 AuthStore.authenticate 验 key，再认 + CODEX_GLOBAL_ADMINS（与 handler_admin._is_global_admin 同款 + 口径：逗号分隔的用户 id 清单，未配置或为空时一律 fail-closed， + 绝不回落到项目角色推导）。任一步失败都抛 PermissionError， + 且在调 reap 之前就拒掉——未授权调用绝不走到判定/杀。 + 被拒的已认证非管理员用户记一条 admin_worker_sweep_denied + 审计（审计写失败不挡拒绝本身）；key 本身验不过时没有可 + 归属的用户，不记这条。 +- 杀前审计 + 执行后复核：授权后直接调 reap_stale_workers， + 复用 classify_stale_workers 的判定，不另造判定函数；每行 + kill 档在 reap 内先再验 pid 已死、杀前写审计（审计失败则 + 不杀，fail-closed）、只给该 pid 本人发信号、杀后用同一 + checker 复核并写杀后审计；本 CLI 另写一条 + admin_worker_sweep 汇总审计（killed / listed / 失败计数， + 写失败只记 audit_summary=False，不改已执行结果）。 +- 不新增 HTTP 路由（v1.1 契约确认 B8 未定前不扩 API 面）； + 不枚举容器、不按名/前缀匹配、不碰基础设施容器、不删 + workers 注册行：杀的是判定为 kill 的行对应的 pid 本人， + 注册行的去留是另一条决策，本 CLI 不擅自删。 +- --dry-run 只调 classify_stale_workers 列判定，不调 reap、 + 不杀、不写杀审计，供运维先看再动手。 + +本模块只 import cloud.auth 与 cloud.workers，不 import +cloud.server / handler_*，避免循环 import；PG 库路径同样经 +cloud.db 的 connect 适配（与 workers 本体同款）。 +""" + +import argparse +import json +import os +import sys + +from cloud.auth import AuthStore +from cloud.workers import classify_stale_workers +from cloud.workers import reap_stale_workers + + +def global_admin_ids(admin_ids=None): + """全局管理员清单：显式传入优先，否则读 CODEX_GLOBAL_ADMINS。 + + 未配置或为空时返回空集——调用方一律 fail-closed。 + """ + if admin_ids is not None: + return {str(item).strip() for item in admin_ids if str(item).strip()} + raw = os.environ.get("CODEX_GLOBAL_ADMINS") or "" + return {part.strip() for part in raw.split(",") if part.strip()} + + +def sweep_stale_workers_admin(db_path, api_key, stale_seconds=None, + local_host=None, pid_is_dead_fn=None, + kill_fn=None, audit_fn=None, dry_run=False, + admin_ids=None, now=None): + """管理员清扫入口：鉴权通过后才调 reap（或 dry-run 只列判定）。 + + 返回 dict：user_id、dry_run、results（每行判定 + 执行字段， + 与 reap/classify 的返回同款）、killed / listed / failed + 三个 worker_id 清单、audit_summary（汇总审计是否写成功， + dry-run 时为 None）。未授权时抛 PermissionError 且绝不调 + reap、不调 kill_fn。kill_fn / audit_fn / pid_is_dead_fn 仅 + 供测试注入；生产调用一律走默认（真 SIGKILL + 统一审计口）。 + """ + store = AuthStore(db_path) + try: + user_id = store.authenticate(api_key) if api_key else None + if user_id is None: + raise PermissionError( + "worker sweep requires a valid admin API key") + if user_id not in global_admin_ids(admin_ids): + try: + store.audit(user_id, "", "admin_worker_sweep_denied", + "worker sweep forbidden: not a global admin") + except Exception: + pass + raise PermissionError( + "worker sweep forbidden: not a global admin") + if dry_run: + results = classify_stale_workers( + db_path, stale_seconds=stale_seconds, + local_host=local_host, pid_is_dead_fn=pid_is_dead_fn, + now=now) + return { + "user_id": user_id, + "dry_run": True, + "results": results, + "killed": [], + "listed": [row["worker_id"] for row in results], + "failed": [], + "audit_summary": None, + } + results = reap_stale_workers( + db_path, stale_seconds=stale_seconds, local_host=local_host, + pid_is_dead_fn=pid_is_dead_fn, kill_fn=kill_fn, + audit_fn=audit_fn, now=now) + killed = [row["worker_id"] for row in results + if row.get("outcome") == "killed"] + listed = [row["worker_id"] for row in results + if row.get("outcome") == "listed"] + failed = [row["worker_id"] for row in results + if row.get("outcome") in ( + "kill_failed_pid_alive", "audit_failed")] + detail = ("killed=%d listed=%d failed=%d killed_ids=%s" + % (len(killed), len(listed), len(failed), + ",".join(killed))) + try: + store.audit(user_id, "", "admin_worker_sweep", detail) + audit_summary = True + except Exception: + audit_summary = False + return { + "user_id": user_id, + "dry_run": False, + "results": results, + "killed": killed, + "listed": listed, + "failed": failed, + "audit_summary": audit_summary, + } + finally: + try: + store.conn.close() + except Exception: + pass + + +def main(argv=None): + """CLI 入口：鉴权失败 exit 1，参数缺失 exit 2，成功 exit 0。 + + 成功时 stdout 只打 JSON 汇总（worker_id 与 outcome，不含 + 任何密钥）；失败原因打 stderr，不回 JSON。 + """ + parser = argparse.ArgumentParser( + description="管理员清扫失联 worker（只杀 pid 已死验证过的本机行）") + parser.add_argument("--db", required=True, help="控制面数据库路径") + parser.add_argument("--api-key", default=None, + help="管理员 API key；不给时读环境变量 " + "CLOUD_WORKER_SWEEP_API_KEY") + parser.add_argument("--stale-seconds", type=float, default=None, + help="失联阈值秒数；不给时复用 alive 阈值语义") + parser.add_argument("--host", default=None, + help="本机主机名；不给时取 socket.gethostname()") + parser.add_argument("--dry-run", action="store_true", + help="只列判定，不杀") + args = parser.parse_args(argv) + api_key = args.api_key or os.environ.get("CLOUD_WORKER_SWEEP_API_KEY") + try: + summary = sweep_stale_workers_admin( + args.db, api_key, stale_seconds=args.stale_seconds, + local_host=args.host, dry_run=args.dry_run) + except PermissionError as exc: + print("worker sweep refused: %s" % exc, file=sys.stderr) + raise SystemExit(1) + print(json.dumps({ + "user_id": summary["user_id"], + "dry_run": summary["dry_run"], + "killed": summary["killed"], + "listed": summary["listed"], + "failed": summary["failed"], + "audit_summary": summary["audit_summary"], + }, ensure_ascii=False)) + return 0 + + +if __name__ == "__main__": + main() diff --git a/cloud/workers.py b/cloud/workers.py index 6dc1782..681d39b 100644 --- a/cloud/workers.py +++ b/cloud/workers.py @@ -58,6 +58,9 @@ 注册行的去留是另一条决策（删行会改 /workers 列表语义）， 本轨不擅自删。亦不新增 HTTP 路由（会改 v1.1 契约）：调 用方（管理员/清扫器）在有权限的上下文里直接调本函数。 + 生产调用方（A5 接线）：cloud/worker_sweep.py 的管理员 CLI + （鉴权 + reap 全链路），见该模块 docstring；除它与测试外 + 不再有其它生产调用方。 """ import os