#!/usr/bin/env python3
"""Agent for hunting Windows persistence mechanisms across registry, services, and scheduled tasks."""

import json
import argparse
import subprocess
import re
from datetime import datetime

REGISTRY_PERSISTENCE_KEYS = [
    r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\Run",
    r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce",
    r"HKCU\SOFTWARE\Microsoft\Windows\CurrentVersion\Run",
    r"HKCU\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce",
    r"HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon",
    r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\Explorer\Run",
    r"HKLM\SOFTWARE\Microsoft\Active Setup\Installed Components",
    r"HKLM\SYSTEM\CurrentControlSet\Services",
    r"HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options",
    r"HKCU\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\User Shell Folders",
    r"HKCU\Environment",
]

SUSPICIOUS_INDICATORS = [
    r"\\temp\\", r"\\tmp\\", r"\\appdata\\local\\temp\\",
    r"powershell.*-enc", r"cmd\.exe.*/c\s+",
    r"\\users\\public\\", r"\\programdata\\",
    r"mshta\.exe", r"rundll32\.exe", r"regsvr32\.exe",
    r"wscript\.exe", r"cscript\.exe",
    r"base64", r"iex\s*\(", r"downloadstring",
]


def enumerate_registry_persistence():
    """Enumerate common Windows registry persistence locations using reg query."""
    findings = []
    for key in REGISTRY_PERSISTENCE_KEYS:
        try:
            result = subprocess.run(
                ["reg", "query", key], capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                entries = parse_reg_output(result.stdout, key)
                for entry in entries:
                    entry["suspicious"] = any(
                        re.search(p, entry.get("value", ""), re.I)
                        for p in SUSPICIOUS_INDICATORS
                    )
                    findings.append(entry)
        except (subprocess.TimeoutExpired, FileNotFoundError):
            continue
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "total_entries": len(findings),
        "suspicious_entries": sum(1 for f in findings if f.get("suspicious")),
        "findings": findings,
    }


def parse_reg_output(output, parent_key):
    """Parse reg query output into structured entries."""
    entries = []
    current_key = parent_key
    for line in output.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        if line.startswith("HK"):
            current_key = line
            continue
        parts = re.split(r"\s{2,}", line, maxsplit=2)
        if len(parts) >= 3:
            entries.append({
                "key": current_key,
                "name": parts[0],
                "type": parts[1],
                "value": parts[2],
            })
    return entries


def enumerate_scheduled_tasks():
    """List scheduled tasks and flag suspicious ones."""
    try:
        result = subprocess.run(
            ["schtasks", "/query", "/fo", "CSV", "/v"],
            capture_output=True, text=True, timeout=30
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return {"error": "schtasks not available"}
    findings = []
    import csv
    from io import StringIO
    reader = csv.DictReader(StringIO(result.stdout))
    for row in reader:
        task_name = row.get("TaskName", "")
        action = row.get("Task To Run", "")
        author = row.get("Author", "")
        suspicious = any(re.search(p, action, re.I) for p in SUSPICIOUS_INDICATORS)
        if suspicious or "\\Microsoft\\" not in task_name:
            findings.append({
                "task_name": task_name,
                "action": action[:500],
                "author": author,
                "status": row.get("Status", ""),
                "next_run": row.get("Next Run Time", ""),
                "suspicious": suspicious,
            })
    return {
        "total_tasks": len(findings),
        "suspicious_tasks": sum(1 for f in findings if f["suspicious"]),
        "findings": findings,
    }


def enumerate_services():
    """List Windows services and flag those running from unusual paths."""
    try:
        result = subprocess.run(
            ["wmic", "service", "get", "Name,PathName,StartMode,State", "/format:csv"],
            capture_output=True, text=True, timeout=30
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return {"error": "wmic not available"}
    findings = []
    for line in result.stdout.strip().split("\n")[1:]:
        parts = line.strip().split(",")
        if len(parts) >= 5:
            path = parts[3]
            suspicious = any(re.search(p, path, re.I) for p in SUSPICIOUS_INDICATORS)
            findings.append({
                "name": parts[1], "path": path,
                "start_mode": parts[4] if len(parts) > 4 else "",
                "state": parts[2], "suspicious": suspicious,
            })
    return {
        "total_services": len(findings),
        "suspicious_services": sum(1 for f in findings if f["suspicious"]),
        "findings": [f for f in findings if f["suspicious"]],
    }


def main():
    parser = argparse.ArgumentParser(description="Hunt for Windows persistence mechanisms")
    sub = parser.add_subparsers(dest="command")
    sub.add_parser("registry", help="Enumerate registry persistence keys")
    sub.add_parser("tasks", help="Enumerate scheduled tasks")
    sub.add_parser("services", help="Enumerate suspicious services")
    sub.add_parser("all", help="Run all persistence hunts")
    args = parser.parse_args()
    if args.command == "registry":
        result = enumerate_registry_persistence()
    elif args.command == "tasks":
        result = enumerate_scheduled_tasks()
    elif args.command == "services":
        result = enumerate_services()
    elif args.command == "all":
        result = {
            "registry": enumerate_registry_persistence(),
            "scheduled_tasks": enumerate_scheduled_tasks(),
            "services": enumerate_services(),
            "timestamp": datetime.utcnow().isoformat(),
        }
    else:
        parser.print_help()
        return
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()