AirCoding/AirPlan/docs/spec/AirPlan-ParaV2/plugins/airndb/scripts/airndb_capture.py

#!/usr/bin/env python3
"""Run safe bounded tcpdump/WinDump capture helpers for AirNDB."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import shlex
import shutil
import subprocess
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence


DEFAULT_COUNT = 200
DEFAULT_TIMEOUT = 30
MAX_COUNT = 500000
WINDUMP_DOWNLOAD_PAGE = "https://www.winpcap.org/windump/install/"
WINDUMP_DOWNLOAD_URL = "https://www.winpcap.org/windump/install/bin/windump_3_9_5/WinDump.exe"
WINDUMP_SHA1 = "d59bc54721951dec855cbb4bbc000f9a71ea4d95"


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


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def is_windows() -> bool:
    return platform.system().lower().startswith("win")


def parse_env_file(path: Path) -> Dict[str, str]:
    if not path.exists():
        return {}
    values: Dict[str, str] = {}
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key.strip()] = value.strip().strip('"')
    return values


def project_tool_env(project_root: Optional[Path]) -> Dict[str, str]:
    if not project_root:
        return {}
    return parse_env_file(project_root / "AirPlan" / "state" / "airndb" / "tool.env")


def tool_dir(project_root: Optional[Path] = None) -> Path:
    configured = os.environ.get("AIRNDB_TOOL_DIR", "").strip()
    if configured:
        return Path(configured).expanduser().resolve()
    if project_root:
        return (project_root / "AirPlan" / "state" / "airndb" / "tools").resolve()
    return Path.home() / ".airndb" / "tools"


def detect_tool(explicit: str = "", project_root: Optional[Path] = None) -> Optional[str]:
    candidates: List[str] = []
    if explicit:
        candidates.append(explicit)
    project_env = project_tool_env(project_root)
    project_tool = project_env.get("AIRNDB_TCPDUMP", "")
    if project_tool:
        candidates.append(project_tool)
    env_tool = os.environ.get("AIRNDB_TCPDUMP", "")
    if env_tool:
        candidates.append(env_tool)
    bundled = tool_dir(project_root) / "WinDump.exe"
    if bundled.exists():
        candidates.append(str(bundled))
    candidates.extend(["windump", "WinDump.exe", "tcpdump"] if is_windows() else ["tcpdump", "windump", "WinDump.exe"])
    for candidate in candidates:
        if Path(candidate).exists():
            return str(Path(candidate).resolve())
        found = shutil.which(candidate)
        if found:
            return found
    return None


def sha1_file(path: Path) -> str:
    digest = hashlib.sha1()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_project_tool_env(project_root: Path, tool_path: Path) -> Path:
    airndb_dir = project_root / "AirPlan" / "state" / "airndb"
    airndb_dir.mkdir(parents=True, exist_ok=True)
    env_path = airndb_dir / "tool.env"
    env_path.write_text(f"AIRNDB_TCPDUMP={tool_path}\n", encoding="utf-8", newline="\n")
    gitignore = airndb_dir / ".gitignore"
    existing = gitignore.read_text(encoding="utf-8") if gitignore.exists() else ""
    if "tool.env" not in existing.splitlines():
        gitignore.write_text((existing.rstrip() + "\ntool.env\n").lstrip(), encoding="utf-8", newline="\n")
    return env_path


def download_windump(target: Path, timeout: int = 60) -> Dict[str, Any]:
    target.parent.mkdir(parents=True, exist_ok=True)
    temp = target.with_suffix(".download")
    try:
        with urllib.request.urlopen(WINDUMP_DOWNLOAD_URL, timeout=timeout) as response:
            data = response.read()
        temp.write_bytes(data)
        actual_sha1 = sha1_file(temp)
        if actual_sha1.lower() != WINDUMP_SHA1.lower():
            temp.unlink(missing_ok=True)
            return {
                "status": "failed",
                "reason": "sha1_mismatch",
                "expectedSha1": WINDUMP_SHA1,
                "actualSha1": actual_sha1,
                "downloadPage": WINDUMP_DOWNLOAD_PAGE,
                "downloadUrl": WINDUMP_DOWNLOAD_URL,
            }
        if target.exists():
            target.unlink()
        temp.replace(target)
        return {
            "status": "downloaded",
            "path": str(target),
            "sha1": actual_sha1,
            "downloadPage": WINDUMP_DOWNLOAD_PAGE,
            "downloadUrl": WINDUMP_DOWNLOAD_URL,
        }
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        temp.unlink(missing_ok=True)
        return {
            "status": "failed",
            "reason": type(exc).__name__,
            "message": str(exc),
            "downloadPage": WINDUMP_DOWNLOAD_PAGE,
            "downloadUrl": WINDUMP_DOWNLOAD_URL,
        }


def ensure_capture_tool(project_root: Path, explicit: str = "", auto_install: bool = False) -> Dict[str, Any]:
    existing = detect_tool(explicit, project_root)
    if existing:
        env_path = write_project_tool_env(project_root, Path(existing)) if Path(existing).exists() else None
        return {
            "status": "ok",
            "tool": existing,
            "source": "existing",
            "projectEnv": str(env_path) if env_path else "",
        }
    if not is_windows():
        return {
            "status": "missing",
            "tool": "",
            "reason": "tcpdump_or_windump_not_found",
            "hint": "Install tcpdump or pass --tool <path>.",
        }
    if not auto_install:
        return {
            "status": "missing",
            "tool": "",
            "reason": "tcpdump_or_windump_not_found",
            "hint": "Run /airndb enter or pass --tool <path>.",
            "downloadPage": WINDUMP_DOWNLOAD_PAGE,
        }
    target = tool_dir(project_root) / "WinDump.exe"
    download = download_windump(target)
    if download.get("status") != "downloaded":
        return {
            "status": "missing",
            "tool": "",
            "reason": "windump_download_failed",
            "download": download,
            "hint": "Install tcpdump/WinDump manually or set AIRNDB_TCPDUMP.",
        }
    env_path = write_project_tool_env(project_root, target)
    return {
        "status": "ok",
        "tool": str(target),
        "source": "downloaded",
        "projectEnv": str(env_path),
        "download": download,
        "driverHint": "WinDump still requires a packet capture driver such as Npcap or WinPcap; install one manually if interface listing fails.",
    }


def split_filter(filter_expr: str) -> List[str]:
    if not filter_expr.strip():
        return []
    try:
        return shlex.split(filter_expr, posix=not is_windows())
    except ValueError:
        return filter_expr.split()


def display_command(command: Sequence[str]) -> str:
    if is_windows():
        return subprocess.list2cmdline(list(command))
    return shlex.join(list(command))


def tail(text: str, limit: int = 8000) -> str:
    if len(text) <= limit:
        return text
    return text[-limit:]


def bounded_count(value: int) -> int:
    if value < 1:
        raise SystemExit("--count must be >= 1")
    if value > MAX_COUNT:
        raise SystemExit(f"--count must be <= {MAX_COUNT}")
    return value


def capture_dir(project_root: Path, output_dir: str = "") -> Path:
    path = Path(output_dir).expanduser().resolve() if output_dir else project_root / "AirPlan" / "docs" / "network" / "airndb-captures"
    path.mkdir(parents=True, exist_ok=True)
    return path


def base_capture_command(tool: str, iface: str, output: Path, count: int, filter_expr: str) -> List[str]:
    if not iface:
        raise SystemExit("--iface is required for capture or command")
    command = [tool, "-i", iface, "-nn", "-s", "0", "-w", str(output), "-c", str(bounded_count(count))]
    command.extend(split_filter(filter_expr))
    return command


def read_command(tool: str, read_file: Path, filter_expr: str) -> List[str]:
    if not read_file.exists():
        raise SystemExit(f"read file not found: {read_file}")
    command = [tool, "-nn", "-tttt", "-r", str(read_file)]
    command.extend(split_filter(filter_expr))
    return command


def run_command(command: Sequence[str], timeout: int) -> Dict[str, Any]:
    started = now_iso()
    try:
        completed = subprocess.run(
            list(command),
            check=False,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout,
        )
        return {
            "startedAt": started,
            "finishedAt": now_iso(),
            "timedOut": False,
            "returnCode": completed.returncode,
            "stdout": tail(completed.stdout),
            "stderr": tail(completed.stderr),
        }
    except subprocess.TimeoutExpired as exc:
        return {
            "startedAt": started,
            "finishedAt": now_iso(),
            "timedOut": True,
            "returnCode": None,
            "stdout": tail((exc.stdout or "") if isinstance(exc.stdout, str) else ""),
            "stderr": tail((exc.stderr or "") if isinstance(exc.stderr, str) else ""),
        }


def write_report(output_dir: Path, action: str, payload: Dict[str, Any]) -> Path:
    report_path = output_dir / f"{now_stamp()}-airndb-{action}.json"
    report_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return report_path


def append_log(project_root: Path, payload: Dict[str, Any]) -> None:
    log_path = project_root / "AirPlan" / "docs" / "network" / "airndb-log.md"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    lines = [
        "",
        f"## {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}: AirNDB {payload['action']}",
        "",
        f"- Tool: `{payload.get('tool', 'missing')}`",
        f"- Interface: `{payload.get('interface', '')}`",
        f"- Filter: `{payload.get('filter', '')}`",
        f"- Command: `{payload.get('command', '')}`",
        f"- Artifacts: `{payload.get('artifact', '')}`",
        f"- Report: `{payload.get('report', '')}`",
        f"- Result: `{payload.get('result', '')}`",
        "- AirDbg handoff: TODO",
        "- Residual risk: pcap files may contain sensitive packet contents; review before sharing.",
    ]
    with log_path.open("a", encoding="utf-8", newline="\n") as handle:
        handle.write("\n".join(lines) + "\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="AirNDB tcpdump/WinDump helper.")
    parser.add_argument("--project", default=".")
    parser.add_argument("--action", choices=["interfaces", "command", "capture", "read"], default="interfaces")
    parser.add_argument("--tool", default="", help="Explicit tcpdump/WinDump path.")
    parser.add_argument("--iface", default="", help="Interface name or WinDump interface index.")
    parser.add_argument("--filter", default="", help="BPF filter expression.")
    parser.add_argument("--count", type=int, default=DEFAULT_COUNT)
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT)
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--output", default="")
    parser.add_argument("--read-file", default="")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    project_root = Path(args.project).expanduser().resolve()
    output_dir = capture_dir(project_root, args.output_dir)
    tool = detect_tool(args.tool, project_root)
    if not tool:
        print("airndb_status=blocked")
        print("reason=tcpdump_or_windump_not_found")
        print("hint=Install tcpdump on Unix-like systems or WinDump/Npcap on Windows, or pass --tool <path>.")
        raise SystemExit(2)

    if args.action == "interfaces":
        command = [tool, "-D"]
        result = run_command(command, max(5, args.timeout))
        payload = {
            "action": "interfaces",
            "tool": tool,
            "command": display_command(command),
            "result": result,
        }
        report = write_report(output_dir, "interfaces", payload)
        print("airndb_status=ok" if result["returnCode"] == 0 else "airndb_status=failed")
        print(f"tool={tool}")
        print(f"command={display_command(command)}")
        print(f"report={report}")
        print(result.get("stdout", ""))
        print(result.get("stderr", ""))
        append_log(project_root, {
            "action": "interfaces",
            "tool": tool,
            "interface": "",
            "filter": "",
            "command": display_command(command),
            "artifact": "",
            "report": str(report),
            "result": "ok" if result["returnCode"] == 0 else "failed",
        })
        return

    if args.action == "command":
        output = Path(args.output).expanduser().resolve() if args.output else output_dir / f"{now_stamp()}-airndb.pcap"
        command = base_capture_command(tool, args.iface, output, args.count, args.filter)
        print("airndb_status=ok")
        print(f"capture_command={display_command(command)}")
        print(f"output={output}")
        return

    if args.action == "capture":
        output = Path(args.output).expanduser().resolve() if args.output else output_dir / f"{now_stamp()}-airndb.pcap"
        command = base_capture_command(tool, args.iface, output, args.count, args.filter)
        result = run_command(command, max(1, args.timeout))
        payload = {
            "action": "capture",
            "tool": tool,
            "interface": args.iface,
            "filter": args.filter,
            "count": bounded_count(args.count),
            "timeoutSeconds": args.timeout,
            "pcap": str(output),
            "command": display_command(command),
            "result": result,
        }
        report = write_report(output_dir, "capture", payload)
        status = "timeout" if result["timedOut"] else ("ok" if result["returnCode"] == 0 else "failed")
        print(f"airndb_status={status}")
        print(f"pcap={output}")
        print(f"report={report}")
        print(f"command={display_command(command)}")
        append_log(project_root, {
            "action": "capture",
            "tool": tool,
            "interface": args.iface,
            "filter": args.filter,
            "command": display_command(command),
            "artifact": str(output),
            "report": str(report),
            "result": status,
        })
        return

    if args.action == "read":
        if not args.read_file:
            raise SystemExit("--read-file is required for read")
        read_file = Path(args.read_file).expanduser()
        if not read_file.is_absolute():
            read_file = project_root / read_file
        read_file = read_file.resolve()
        command = read_command(tool, read_file, args.filter)
        result = run_command(command, max(1, args.timeout))
        summary_path = output_dir / f"{now_stamp()}-airndb-read.txt"
        summary_path.write_text((result.get("stdout") or "") + (result.get("stderr") or ""), encoding="utf-8", newline="\n")
        payload = {
            "action": "read",
            "tool": tool,
            "readFile": str(read_file),
            "filter": args.filter,
            "summary": str(summary_path),
            "command": display_command(command),
            "result": result,
        }
        report = write_report(output_dir, "read", payload)
        status = "ok" if result["returnCode"] == 0 else "failed"
        print(f"airndb_status={status}")
        print(f"summary={summary_path}")
        print(f"report={report}")
        print(f"command={display_command(command)}")
        append_log(project_root, {
            "action": "read",
            "tool": tool,
            "interface": "",
            "filter": args.filter,
            "command": display_command(command),
            "artifact": str(summary_path),
            "report": str(report),
            "result": status,
        })


if __name__ == "__main__":
    main()