run_python.py

"""
run_python.py
Analyst toolkit -- Python script runner.

Single-file CLI, class-based, mirrors the --mode / --output-json pattern
used across the toolkit (same convention as the security-agent detection
engine's --mode / --seed / --output-json CLI): --mode batch runs every
registered script once and exits; --mode menu is an interactive on-demand
picker for use mid-engagement.

Usage:
    python run_python.py --mode batch
    python run_python.py --mode menu
    python run_python.py --mode batch --output-json results.json
"""

import argparse
import json
import subprocess
import sys
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path
from typing import List, Optional


def timestamp() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


@dataclass
class ScriptResult:
    path: str
    ok: bool
    exit_code: Optional[int]
    started: str
    finished: str
    log_file: str


class PythonScriptRunner:
    """Loads registered Python scripts from scripts_config.json and
    executes them, either as a single batch pass or via an interactive
    on-demand menu."""

    def __init__(self, config_path: Path, log_dir: Path):
        self.config_path = config_path
        self.log_dir = log_dir
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.scripts: List[dict] = self._load_config()
        self.results: List[ScriptResult] = []

    def _load_config(self) -> List[dict]:
        if not self.config_path.exists():
            print(f"[ERROR] Config file not found: {self.config_path}")
            sys.exit(1)
        with open(self.config_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data.get("python", [])

    def _run_one(self, entry: dict, extra_args: Optional[List[str]] = None) -> ScriptResult:
        script_path = entry.get("path")
        args = list(entry.get("args", [])) + list(extra_args or [])
        started = timestamp()

        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_name = Path(script_path).stem if script_path else "unknown"
        log_path = self.log_dir / f"{safe_name}_{stamp}.log"

        if not script_path or not Path(script_path).exists():
            msg = f"[SKIP] Script not found: {script_path}"
            print(msg)
            log_path.write_text(msg + "\n", encoding="utf-8")
            return ScriptResult(script_path or "unknown", False, None, started, timestamp(), str(log_path))

        header = f"=== [{started}] Running: {script_path} {' '.join(args)} ==="
        print("\n" + header)

        try:
            result = subprocess.run(
                [sys.executable, script_path, *args],
                capture_output=True,
                text=True,
                timeout=600,  # safety timeout, adjust per script if needed
            )
            log_path.write_text(
                header + "\n--- STDOUT ---\n" + (result.stdout or "")
                + "\n--- STDERR ---\n" + (result.stderr or ""),
                encoding="utf-8",
            )
            ok = result.returncode == 0
            status = "[OK]" if ok else "[FAIL]"
            print(f"{status} {script_path} (exit {result.returncode})")
            return ScriptResult(script_path, ok, result.returncode, started, timestamp(), str(log_path))

        except subprocess.TimeoutExpired:
            msg = f"[TIMEOUT] {script_path} exceeded time limit"
            print(msg)
            log_path.write_text(header + "\n" + msg, encoding="utf-8")
            return ScriptResult(script_path, False, None, started, timestamp(), str(log_path))
        except Exception as e:
            msg = f"[ERROR] {script_path} raised {e}"
            print(msg)
            log_path.write_text(header + "\n" + msg, encoding="utf-8")
            return ScriptResult(script_path, False, None, started, timestamp(), str(log_path))

    def run_batch(self) -> List[ScriptResult]:
        if not self.scripts:
            print("No python scripts configured in scripts_config.json.")
            return []
        self.results = [self._run_one(entry) for entry in self.scripts]
        self._print_summary()
        return self.results

    def run_menu(self) -> None:
        while True:
            self.scripts = self._load_config()  # reload so newly added scripts show up
            if not self.scripts:
                print("No python scripts configured.")
                return

            print("\n================ Python Script Menu ================")
            for i, entry in enumerate(self.scripts, start=1):
                print(f"  [{i}] {entry.get('path')}")
            print("  [R] Reload config    [Q] Quit")
            print("======================================================")

            choice = input("Pick a script to run: ").strip()
            if choice.lower() == "q":
                break
            if choice.lower() == "r":
                continue
            if choice.isdigit() and 1 <= int(choice) <= len(self.scripts):
                entry = self.scripts[int(choice) - 1]
                extra = input("Extra args for this run (blank for none): ").strip()
                extra_args = extra.split() if extra else []
                self.results.append(self._run_one(entry, extra_args))
            else:
                print("Invalid choice.")

    def _print_summary(self) -> None:
        passed = sum(1 for r in self.results if r.ok)
        total = len(self.results)
        print(f"\n=== Summary: {passed}/{total} scripts succeeded ===")
        for r in self.results:
            print(f"  {'OK' if r.ok else 'FAIL'} - {r.path}")

    def write_json_summary(self, output_path: Path) -> None:
        payload = {
            "generated": timestamp(),
            "total": len(self.results),
            "passed": sum(1 for r in self.results if r.ok),
            "results": [asdict(r) for r in self.results],
        }
        output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"JSON summary written to {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Analyst toolkit -- Python script runner.")
    parser.add_argument("--config", default="scripts_config.json", help="Path to config JSON")
    parser.add_argument("--log-dir", default="logs", help="Directory to write log files")
    parser.add_argument("--mode", choices=["batch", "menu"], default="batch",
                         help="batch = run every registered script once; menu = interactive on-demand")
    parser.add_argument("--output-json", default=None,
                         help="Optional path to write a JSON run summary (batch mode only)")
    args = parser.parse_args()

    runner = PythonScriptRunner(Path(args.config), Path(args.log_dir))

    if args.mode == "batch":
        runner.run_batch()
        if args.output_json:
            runner.write_json_summary(Path(args.output_json))
    else:
        runner.run_menu()


if __name__ == "__main__":
    main()