106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
||
"""Run the xdfx flow one step at a time.
|
||
|
||
The heavy business logic intentionally stays in xdfx_flow.py. This wrapper only
|
||
keeps "full flow" orchestration readable and preserves per-step feedback: every
|
||
subprocess line is flushed immediately, so OP/SAVE/ERROR/STEP feedback remains
|
||
visible to the calling agent.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
if hasattr(sys.stdout, "reconfigure"):
|
||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||
if hasattr(sys.stderr, "reconfigure"):
|
||
sys.stderr.reconfigure(encoding="utf-8", errors="replace")
|
||
|
||
STEPS = ("step1", "step2", "step3", "step4", "step5", "step6")
|
||
|
||
STEP_LABELS = {
|
||
"step1": "受众目标分析与入库",
|
||
"step2": "关键传播节点分析与入库",
|
||
"step3": "受众脆弱点分析与入库",
|
||
"step4": "传播节点脆弱点分析与入库",
|
||
"step5": "受众策略分析与入库",
|
||
"step6": "传播节点策略分析与入库",
|
||
}
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
"""Parse only wrapper-owned flags; pass all other flags to xdfx_flow.py."""
|
||
parser = argparse.ArgumentParser(
|
||
description="按六个步骤运行 xdfx_flow.py，并实时转发中文反馈日志。",
|
||
add_help=True,
|
||
)
|
||
parser.add_argument("--from-step", choices=STEPS, default="step1")
|
||
parser.add_argument("--to-step", choices=STEPS, default="step6")
|
||
parser.add_argument(
|
||
"--continue-on-error",
|
||
action="store_true",
|
||
help="某一步返回非零退出码时仍继续执行后续步骤。",
|
||
)
|
||
return parser
|
||
|
||
|
||
def selected_steps(first: str, last: str) -> list[str]:
|
||
"""Return an inclusive step slice for resume/partial reruns."""
|
||
start = STEPS.index(first)
|
||
end = STEPS.index(last)
|
||
if start > end:
|
||
raise SystemExit("起始步骤必须小于或等于结束步骤")
|
||
return list(STEPS[start : end + 1])
|
||
|
||
|
||
def stream_command(step: str, cmd: list[str]) -> tuple[int, bool]:
|
||
"""Run one step and forward stdout/stderr lines as they are produced."""
|
||
print(f"运行步骤：{STEP_LABELS.get(step, step)}；执行命令：" + " ".join(cmd), flush=True)
|
||
env = os.environ.copy()
|
||
env["PYTHONIOENCODING"] = "utf-8"
|
||
process = subprocess.Popen(
|
||
cmd,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
encoding="utf-8",
|
||
errors="replace",
|
||
env=env,
|
||
)
|
||
assert process.stdout is not None
|
||
waiting_user = False
|
||
for line in process.stdout:
|
||
if line.startswith("用户协助 "):
|
||
waiting_user = True
|
||
sys.stdout.buffer.write(line.encode("utf-8", errors="replace"))
|
||
sys.stdout.buffer.flush()
|
||
return process.wait(), waiting_user
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
parser = build_parser()
|
||
args, passthrough = parser.parse_known_args(argv)
|
||
flow_script = Path(__file__).with_name("xdfx_flow.py")
|
||
exit_code = 0
|
||
for step in selected_steps(args.from_step, args.to_step):
|
||
# Keep all business flags unchanged so the wrapper behaves like a thin
|
||
# streaming version of repeated xdfx_flow.py --mode stepN calls.
|
||
cmd = [sys.executable, str(flow_script), "--mode", step, *passthrough]
|
||
code, waiting_user = stream_command(step, cmd)
|
||
if waiting_user:
|
||
print("已发起用户协助，暂停后续步骤；请根据用户选择重跑当前步骤。", flush=True)
|
||
return code
|
||
if code:
|
||
exit_code = code
|
||
if not args.continue_on_error:
|
||
return code
|
||
return exit_code
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|