Anthropic-Cybersecurity-Skills/skills/eradicating-malware-from-infected-systems/scripts/agent.py

#!/usr/bin/env python3
"""Agent for malware eradication from infected systems — process kill, file removal, persistence cleanup."""

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime, timezone


PERSISTENCE_LOCATIONS_WINDOWS = [
    r"HKCU\Software\Microsoft\Windows\CurrentVersion\Run",
    r"HKLM\Software\Microsoft\Windows\CurrentVersion\Run",
    r"HKCU\Software\Microsoft\Windows\CurrentVersion\RunOnce",
    r"HKLM\Software\Microsoft\Windows\CurrentVersion\RunOnce",
    r"HKLM\SYSTEM\CurrentControlSet\Services",
]

PERSISTENCE_LOCATIONS_LINUX = [
    "/etc/cron.d", "/var/spool/cron", "/etc/crontab",
    "/etc/init.d", "/etc/systemd/system", "/usr/lib/systemd/system",
    "/etc/rc.local", "~/.bashrc", "~/.bash_profile",
]


def kill_malicious_process(pid):
    """Terminate a malicious process by PID."""
    if sys.platform == "win32":
        cmd = ["taskkill", "/F", "/PID", str(pid)]
    else:
        cmd = ["kill", "-9", str(pid)]
    try:
        subprocess.check_output(cmd, text=True, errors="replace", timeout=10)
        return {"pid": pid, "status": "killed"}
    except subprocess.SubprocessError as e:
        return {"pid": pid, "status": "failed", "error": str(e)}


def quarantine_file(file_path, quarantine_dir):
    """Move malicious file to quarantine directory with metadata."""
    os.makedirs(quarantine_dir, exist_ok=True)
    if not os.path.isfile(file_path):
        return {"file": file_path, "status": "not_found"}
    import hashlib
    with open(file_path, "rb") as f:
        sha256 = hashlib.sha256(f.read()).hexdigest()
    dest = os.path.join(quarantine_dir, f"{sha256}_{os.path.basename(file_path)}.quarantine")
    try:
        os.rename(file_path, dest)
        meta = {
            "original_path": file_path,
            "sha256": sha256,
            "quarantined_at": datetime.now(timezone.utc).isoformat(),
            "quarantine_path": dest,
        }
        with open(dest + ".json", "w") as f:
            json.dump(meta, f, indent=2)
        return {"file": file_path, "status": "quarantined", "sha256": sha256}
    except OSError as e:
        return {"file": file_path, "status": "failed", "error": str(e)}


def scan_persistence_windows():
    """Scan Windows persistence locations for suspicious entries."""
    findings = []
    for key in PERSISTENCE_LOCATIONS_WINDOWS:
        try:
            result = subprocess.check_output(
                ["reg", "query", key], text=True, errors="replace", timeout=5
            )
            for line in result.splitlines():
                line = line.strip()
                if "REG_SZ" in line or "REG_EXPAND_SZ" in line:
                    parts = line.split(None, 2)
                    if len(parts) >= 3:
                        value = parts[2]
                        if any(d in value.lower() for d in ["temp", "appdata\\local\\temp", "public"]):
                            findings.append({
                                "key": key,
                                "name": parts[0],
                                "value": value,
                                "suspicious": True,
                            })
        except subprocess.SubprocessError:
            pass
    return findings


def scan_persistence_linux():
    """Scan Linux persistence locations for suspicious entries."""
    findings = []
    for loc in PERSISTENCE_LOCATIONS_LINUX:
        expanded = os.path.expanduser(loc)
        if os.path.isdir(expanded):
            for f in os.listdir(expanded):
                fpath = os.path.join(expanded, f)
                if os.path.isfile(fpath):
                    try:
                        stat = os.stat(fpath)
                        if stat.st_mtime > (datetime.now().timestamp() - 86400 * 7):
                            findings.append({
                                "path": fpath,
                                "modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
                                "note": "Recently modified persistence file",
                            })
                    except OSError:
                        pass
        elif os.path.isfile(expanded):
            try:
                with open(expanded, "r", errors="replace") as fh:
                    content = fh.read()
                for line in content.splitlines():
                    if any(s in line.lower() for s in ["curl", "wget", "nc ", "ncat", "bash -i"]):
                        findings.append({
                            "file": expanded,
                            "line": line.strip()[:200],
                            "note": "Suspicious command in startup file",
                        })
            except PermissionError:
                pass
    return findings


def clean_scheduled_tasks():
    """List suspicious scheduled tasks."""
    findings = []
    if sys.platform == "win32":
        try:
            result = subprocess.check_output(
                ["schtasks", "/Query", "/FO", "CSV", "/NH", "/V"],
                text=True, errors="replace", timeout=15
            )
            for line in result.splitlines():
                parts = line.strip('"').split('","')
                if len(parts) >= 9:
                    action = parts[8] if len(parts) > 8 else ""
                    if any(s in action.lower() for s in ["powershell", "cmd", "temp", "appdata"]):
                        findings.append({
                            "task_name": parts[1] if len(parts) > 1 else "",
                            "action": action[:200],
                        })
        except subprocess.SubprocessError:
            pass
    return findings


def main():
    parser = argparse.ArgumentParser(
        description="Eradicate malware from infected systems"
    )
    parser.add_argument("--scan-persistence", action="store_true")
    parser.add_argument("--kill-pid", type=int, nargs="*", help="PIDs to terminate")
    parser.add_argument("--quarantine-file", nargs="*", help="Files to quarantine")
    parser.add_argument("--quarantine-dir", default=os.environ.get("QUARANTINE_DIR", "/tmp/quarantine"))
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Malware Eradication Agent")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "actions": []}

    if args.kill_pid:
        for pid in args.kill_pid:
            result = kill_malicious_process(pid)
            report["actions"].append({"action": "kill", **result})
            print(f"[*] Kill PID {pid}: {result['status']}")

    if args.quarantine_file:
        for fpath in args.quarantine_file:
            result = quarantine_file(fpath, args.quarantine_dir)
            report["actions"].append({"action": "quarantine", **result})
            print(f"[*] Quarantine {fpath}: {result['status']}")

    if args.scan_persistence:
        if sys.platform == "win32":
            pers = scan_persistence_windows()
            tasks = clean_scheduled_tasks()
        else:
            pers = scan_persistence_linux()
            tasks = []
        report["persistence"] = pers
        report["scheduled_tasks"] = tasks
        print(f"[*] Persistence entries: {len(pers)}, Suspicious tasks: {len(tasks)}")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"[*] Report saved to {args.output}")
    else:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()