#!/usr/bin/env python3
"""Agent for performing malware persistence investigation.

Enumerates Windows registry Run keys, services, scheduled tasks,
WMI subscriptions, and Linux cron/systemd persistence mechanisms.
"""

import json
import sys
import xml.etree.ElementTree as ET
from pathlib import Path


SUSPICIOUS_INDICATORS = [
    "powershell", "cmd /c", "wscript", "cscript", "mshta",
    "regsvr32", "rundll32", "certutil", "bitsadmin",
    "downloadstring", "invoke-", "iex", "hidden",
    "bypass", "base64", "-enc", "-e ", ".ps1", ".vbs", ".hta",
    "programdata", "appdata\\local\\temp", "/tmp/",
]


class PersistenceInvestigator:
    """Investigates malware persistence mechanisms on forensic images."""

    def __init__(self, evidence_root, output_dir):
        self.evidence_root = Path(evidence_root)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.findings = []

    def _is_suspicious(self, value):
        """Check if a value matches suspicious indicators."""
        value_lower = str(value).lower()
        return any(ind in value_lower for ind in SUSPICIOUS_INDICATORS)

    def check_registry_run_keys(self):
        """Check registry Run keys using python-registry."""
        try:
            from Registry import Registry
        except ImportError:
            return [{"error": "Install python-registry: pip install python-registry"}]

        entries = []
        sw_path = self.evidence_root / "Windows/System32/config/SOFTWARE"
        if not sw_path.exists():
            return entries

        reg = Registry.Registry(str(sw_path))
        run_paths = [
            "Microsoft\\Windows\\CurrentVersion\\Run",
            "Microsoft\\Windows\\CurrentVersion\\RunOnce",
            "Wow6432Node\\Microsoft\\Windows\\CurrentVersion\\Run",
        ]
        for key_path in run_paths:
            try:
                key = reg.open(key_path)
                for value in key.values():
                    entry = {
                        "location": f"HKLM\\SOFTWARE\\{key_path}",
                        "name": value.name(),
                        "value": str(value.value()),
                        "timestamp": str(key.timestamp()),
                        "suspicious": self._is_suspicious(str(value.value())),
                    }
                    entries.append(entry)
                    if entry["suspicious"]:
                        self.findings.append(entry)
            except Exception:
                continue
        return entries

    def check_services(self):
        """Check Windows services for suspicious auto-start entries."""
        try:
            from Registry import Registry
        except ImportError:
            return []

        entries = []
        system_path = self.evidence_root / "Windows/System32/config/SYSTEM"
        if not system_path.exists():
            return entries

        reg = Registry.Registry(str(system_path))
        try:
            select = reg.open("Select")
            current = select.value("Current").value()
            cs = f"ControlSet{current:03d}"
            services = reg.open(f"{cs}\\Services")

            for svc in services.subkeys():
                try:
                    start = svc.value("Start").value()
                    if start not in (0, 1, 2):
                        continue
                    image_path = ""
                    try:
                        image_path = svc.value("ImagePath").value()
                    except Exception:
                        continue
                    entry = {
                        "service_name": svc.name(),
                        "image_path": image_path,
                        "start_type": start,
                        "timestamp": str(svc.timestamp()),
                        "suspicious": self._is_suspicious(image_path),
                    }
                    entries.append(entry)
                    if entry["suspicious"]:
                        self.findings.append(entry)
                except Exception:
                    continue
        except Exception:
            pass
        return entries

    def check_scheduled_tasks(self):
        """Parse Windows scheduled task XML files."""
        tasks_dir = self.evidence_root / "Windows/System32/Tasks"
        if not tasks_dir.exists():
            return []

        entries = []
        ns = {"t": "http://schemas.microsoft.com/windows/2004/02/mit/task"}

        for task_file in tasks_dir.rglob("*"):
            if not task_file.is_file():
                continue
            try:
                tree = ET.parse(str(task_file))
                root = tree.getroot()
                for action in root.findall(".//t:Exec", ns):
                    cmd = action.find("t:Command", ns)
                    args = action.find("t:Arguments", ns)
                    cmd_text = cmd.text if cmd is not None else ""
                    args_text = args.text if args is not None else ""
                    combined = f"{cmd_text} {args_text}"
                    entry = {
                        "task_name": task_file.name,
                        "command": cmd_text,
                        "arguments": args_text,
                        "suspicious": self._is_suspicious(combined),
                    }
                    entries.append(entry)
                    if entry["suspicious"]:
                        self.findings.append(entry)
            except Exception:
                continue
        return entries

    def check_startup_folder(self):
        """Check Windows startup folder contents."""
        entries = []
        startup_paths = [
            "ProgramData/Microsoft/Windows/Start Menu/Programs/Startup",
        ]
        for user_dir in (self.evidence_root / "Users").iterdir() if (self.evidence_root / "Users").exists() else []:
            if user_dir.is_dir():
                startup_paths.append(
                    f"Users/{user_dir.name}/AppData/Roaming/Microsoft/Windows/Start Menu/Programs/Startup"
                )

        for spath in startup_paths:
            target = self.evidence_root / spath
            if target.exists() and target.is_dir():
                for item in target.iterdir():
                    if item.is_file():
                        entries.append({
                            "location": str(spath),
                            "filename": item.name,
                            "size": item.stat().st_size,
                            "suspicious": self._is_suspicious(item.name),
                        })
        return entries

    def check_linux_persistence(self):
        """Check Linux persistence mechanisms."""
        entries = []
        checks = {
            "crontab": "etc/crontab",
            "rc_local": "etc/rc.local",
            "ld_preload": "etc/ld.so.preload",
        }
        for name, path in checks.items():
            target = self.evidence_root / path
            if target.exists() and target.is_file():
                content = target.read_text(errors="ignore")
                entries.append({
                    "mechanism": name,
                    "file": path,
                    "content_preview": content[:500],
                    "suspicious": self._is_suspicious(content),
                })

        cron_dirs = ["etc/cron.d", "etc/cron.daily", "var/spool/cron/crontabs"]
        for cron_dir in cron_dirs:
            target = self.evidence_root / cron_dir
            if target.exists() and target.is_dir():
                for f in target.iterdir():
                    if f.is_file():
                        content = f.read_text(errors="ignore")
                        entries.append({
                            "mechanism": "cron_job",
                            "file": str(f.relative_to(self.evidence_root)),
                            "content_preview": content[:300],
                        })

        auth_keys_pattern = "home/*/. ssh/authorized_keys".replace(". ", ".")
        for ak in self.evidence_root.glob("home/*/.ssh/authorized_keys"):
            entries.append({
                "mechanism": "ssh_authorized_keys",
                "file": str(ak.relative_to(self.evidence_root)),
                "content_preview": ak.read_text(errors="ignore")[:300],
            })
        return entries

    def generate_report(self):
        """Run all persistence checks and generate report."""
        report = {
            "case_output": str(self.output_dir),
            "evidence_root": str(self.evidence_root),
            "registry_run_keys": self.check_registry_run_keys(),
            "services": self.check_services()[:30],
            "scheduled_tasks": self.check_scheduled_tasks(),
            "startup_folder": self.check_startup_folder(),
            "linux_persistence": self.check_linux_persistence(),
            "suspicious_findings": self.findings,
            "total_suspicious": len(self.findings),
        }

        report_path = self.output_dir / "persistence_report.json"
        with open(report_path, "w") as f:
            json.dump(report, f, indent=2)
        print(json.dumps(report, indent=2))
        return report


def main():
    if len(sys.argv) < 3:
        print("Usage: agent.py <evidence_root> <output_dir>")
        sys.exit(1)
    investigator = PersistenceInvestigator(sys.argv[1], sys.argv[2])
    investigator.generate_report()


if __name__ == "__main__":
    main()
