#!/usr/bin/env python3
"""Velociraptor incident response collection agent.

Interfaces with the Velociraptor DFIR platform API to schedule artifact
collections on endpoints, retrieve results, and generate IR reports.
Supports common forensic artifacts including process listings, network
connections, autoruns, event logs, and file system evidence.
"""
import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone

try:
    import requests
except ImportError:
    print("[!] 'requests' library required: pip install requests", file=sys.stderr)
    sys.exit(1)

try:
    import grpc
    HAS_GRPC = True
except ImportError:
    HAS_GRPC = False


def get_velo_config():
    """Return Velociraptor API configuration."""
    api_url = os.environ.get("VELOCIRAPTOR_API_URL", "https://localhost:8001")
    api_key = os.environ.get("VELOCIRAPTOR_API_KEY", "")
    cert_path = os.environ.get("VELOCIRAPTOR_CERT", "")
    return api_url.rstrip("/"), api_key, cert_path


def velo_api_call(api_url, api_key, endpoint, method="GET", data=None, cert_path=None):
    """Make an authenticated API call to Velociraptor."""
    url = f"{api_url}{endpoint}"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    verify = cert_path if cert_path else False
    if method == "POST":
        resp = requests.post(url, headers=headers, json=data, verify=verify, timeout=60)
    else:
        resp = requests.get(url, headers=headers, verify=verify, timeout=60)
    resp.raise_for_status()
    return resp.json()


def search_clients(api_url, api_key, query, cert_path=None):
    """Search for Velociraptor clients by hostname, label, or client ID."""
    print(f"[*] Searching for clients: {query}")
    data = {"query": query, "count": 100}
    result = velo_api_call(api_url, api_key, "/api/v1/SearchClients", "POST", data, cert_path)
    clients = result.get("items", [])
    print(f"[+] Found {len(clients)} client(s)")
    for c in clients:
        os_info = c.get("os_info", {})
        print(f"    {c.get('client_id', 'N/A'):20s} | "
              f"{os_info.get('hostname', 'unknown'):20s} | "
              f"{os_info.get('system', 'unknown'):10s} | "
              f"Last seen: {c.get('last_seen_at', 'N/A')}")
    return clients


def collect_artifact(api_url, api_key, client_id, artifacts, parameters=None, cert_path=None):
    """Schedule an artifact collection on a specific client."""
    print(f"[*] Scheduling collection on {client_id}: {', '.join(artifacts)}")
    specs = []
    for artifact in artifacts:
        spec = {"artifact": artifact}
        if parameters and artifact in parameters:
            spec["parameters"] = {"env": [
                {"key": k, "value": v} for k, v in parameters[artifact].items()
            ]}
        specs.append(spec)

    data = {
        "client_id": client_id,
        "artifacts": artifacts,
        "specs": specs,
    }
    result = velo_api_call(api_url, api_key, "/api/v1/CollectArtifact", "POST", data, cert_path)
    flow_id = result.get("flow_id", "")
    print(f"[+] Collection started, flow ID: {flow_id}")
    return flow_id


def get_flow_status(api_url, api_key, client_id, flow_id, cert_path=None):
    """Check the status of a collection flow."""
    data = {"client_id": client_id, "flow_id": flow_id}
    result = velo_api_call(api_url, api_key, "/api/v1/GetFlowDetails", "POST", data, cert_path)
    context = result.get("context", {})
    state = context.get("state", "UNSET")
    return state, context


def wait_for_collection(api_url, api_key, client_id, flow_id, max_wait=300, cert_path=None):
    """Poll until a collection flow completes."""
    print(f"[*] Waiting for collection to complete (max {max_wait}s)...")
    elapsed = 0
    interval = 10
    while elapsed < max_wait:
        state, context = get_flow_status(api_url, api_key, client_id, flow_id, cert_path)
        if state == "FINISHED":
            total_rows = context.get("total_collected_rows", 0)
            total_bytes = context.get("total_uploaded_bytes", 0)
            print(f"[+] Collection complete: {total_rows} rows, "
                  f"{total_bytes / 1024:.1f} KB uploaded")
            return True, context
        if state == "ERROR":
            print(f"[!] Collection failed: {context.get('status', 'unknown')}", file=sys.stderr)
            return False, context
        print(f"    State: {state} ({elapsed}s elapsed)")
        time.sleep(interval)
        elapsed += interval
    print("[!] Timed out waiting for collection", file=sys.stderr)
    return False, {}


def get_flow_results(api_url, api_key, client_id, flow_id, artifact, cert_path=None):
    """Retrieve collected artifact results."""
    print(f"[*] Retrieving results for {artifact}")
    data = {
        "client_id": client_id,
        "flow_id": flow_id,
        "artifact": artifact,
        "count": 10000,
    }
    result = velo_api_call(api_url, api_key, "/api/v1/GetTable", "POST", data, cert_path)
    rows = result.get("rows", [])
    columns = result.get("columns", [])
    print(f"[+] Retrieved {len(rows)} row(s), {len(columns)} column(s)")
    return rows, columns


IR_ARTIFACT_SETS = {
    "triage": [
        "Windows.System.Pslist",
        "Windows.Network.Netstat",
        "Windows.Sys.Users",
        "Windows.System.TaskScheduler",
        "Generic.System.Pstree",
    ],
    "persistence": [
        "Windows.Sysinternals.Autoruns",
        "Windows.System.TaskScheduler",
        "Windows.Registry.Run",
        "Windows.System.Services",
    ],
    "network": [
        "Windows.Network.Netstat",
        "Windows.Network.ArpCache",
        "Windows.Network.DNSCache",
        "Windows.Network.InterfaceAddresses",
    ],
    "logs": [
        "Windows.EventLogs.Cleared",
        "Windows.EventLogs.RDPAuth",
        "Windows.EventLogs.PowershellScriptblock",
    ],
    "linux_triage": [
        "Linux.Sys.Pslist",
        "Linux.Network.Netstat",
        "Linux.Sys.Users",
        "Linux.Sys.Crontab",
        "Linux.Sys.LastUserLogin",
    ],
}


def format_summary(client_id, artifacts, flow_context, results_summary):
    """Print collection summary."""
    print(f"\n{'='*60}")
    print(f"  Velociraptor IR Collection Report")
    print(f"{'='*60}")
    print(f"  Client ID   : {client_id}")
    print(f"  Flow ID     : {flow_context.get('session_id', 'N/A')}")
    print(f"  State       : {flow_context.get('state', 'N/A')}")
    print(f"  Artifacts   : {len(artifacts)}")
    print(f"  Total Rows  : {flow_context.get('total_collected_rows', 0)}")

    for artifact_name, row_count in results_summary.items():
        print(f"    {artifact_name:50s}: {row_count} rows")


def main():
    parser = argparse.ArgumentParser(
        description="Velociraptor IR collection agent"
    )
    sub = parser.add_subparsers(dest="command")

    p_search = sub.add_parser("search", help="Search for clients")
    p_search.add_argument("--query", required=True, help="Search query (hostname, label, client ID)")

    p_collect = sub.add_parser("collect", help="Collect artifacts from client")
    p_collect.add_argument("--client-id", required=True, help="Velociraptor client ID")
    p_collect.add_argument("--artifacts", nargs="+", help="Specific artifact names to collect")
    p_collect.add_argument("--preset", choices=list(IR_ARTIFACT_SETS.keys()),
                          help="Use a predefined IR artifact set")
    p_collect.add_argument("--wait", type=int, default=300, help="Max wait time in seconds")

    p_status = sub.add_parser("status", help="Check collection status")
    p_status.add_argument("--client-id", required=True)
    p_status.add_argument("--flow-id", required=True)

    parser.add_argument("--api-url", help="Velociraptor API URL (or VELOCIRAPTOR_API_URL env)")
    parser.add_argument("--api-key", help="API key (or VELOCIRAPTOR_API_KEY env)")
    parser.add_argument("--cert", help="CA cert path (or VELOCIRAPTOR_CERT env)")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    if args.api_url:
        os.environ["VELOCIRAPTOR_API_URL"] = args.api_url
    if args.api_key:
        os.environ["VELOCIRAPTOR_API_KEY"] = args.api_key
    if args.cert:
        os.environ["VELOCIRAPTOR_CERT"] = args.cert

    api_url, api_key, cert_path = get_velo_config()
    if not api_key:
        print("[!] Set VELOCIRAPTOR_API_KEY env var or use --api-key", file=sys.stderr)
        sys.exit(1)

    result = {}

    if args.command == "search":
        clients = search_clients(api_url, api_key, args.query, cert_path)
        result = {"action": "search", "query": args.query, "clients": clients}

    elif args.command == "collect":
        artifacts = args.artifacts or IR_ARTIFACT_SETS.get(args.preset, [])
        if not artifacts:
            print("[!] Specify --artifacts or --preset", file=sys.stderr)
            sys.exit(1)
        flow_id = collect_artifact(api_url, api_key, args.client_id, artifacts, cert_path=cert_path)
        success, context = wait_for_collection(api_url, api_key, args.client_id, flow_id, args.wait, cert_path)
        results_summary = {}
        if success:
            for artifact in artifacts:
                try:
                    rows, cols = get_flow_results(api_url, api_key, args.client_id, flow_id, artifact, cert_path)
                    results_summary[artifact] = len(rows)
                except Exception as e:
                    results_summary[artifact] = f"Error: {e}"
        format_summary(args.client_id, artifacts, context, results_summary)
        result = {"action": "collect", "client_id": args.client_id, "flow_id": flow_id,
                  "state": context.get("state", "UNKNOWN"), "artifacts": results_summary}

    elif args.command == "status":
        state, context = get_flow_status(api_url, api_key, args.client_id, args.flow_id, cert_path)
        print(f"[*] Flow {args.flow_id}: {state}")
        result = {"action": "status", "flow_id": args.flow_id, "state": state, "context": context}

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "Velociraptor",
        "result": result,
    }
    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()

"""
Velociraptor IR Collection Automation Script
Manages artifact collection, hunt creation, and result analysis via Velociraptor API.
"""

import json
import os
import csv
import hashlib
from datetime import datetime
from pathlib import Path
from collections import defaultdict


class VelociraptorCollector:
    """Manages Velociraptor artifact collection for incident response."""

    TRIAGE_ARTIFACTS_WINDOWS = [
        "Windows.EventLogs.EvtxHunter",
        "Windows.Forensics.Prefetch",
        "Windows.Registry.AppCompatCache",
        "Windows.Forensics.Amcache",
        "Windows.Forensics.UserAssist",
        "Windows.System.TaskScheduler",
        "Windows.System.Pslist",
        "Windows.Network.Netstat",
        "Windows.Network.DNSCache",
        "Windows.Forensics.PowerShellHistory",
        "Windows.System.Services",
        "Windows.System.StartupItems",
        "Windows.Persistence.PermanentWMIEvents",
    ]

    TRIAGE_ARTIFACTS_LINUX = [
        "Linux.Sys.AuthLogs",
        "Linux.Forensics.BashHistory",
        "Linux.Sys.Crontab",
        "Linux.Sys.Pslist",
        "Linux.Network.Netstat",
        "Linux.Ssh.AuthorizedKeys",
        "Linux.Services",
    ]

    def __init__(self, server_url=None, api_key=None, output_dir="velociraptor_results"):
        self.server_url = server_url or "https://localhost:8001"
        self.api_key = api_key
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.collection_manifest = []

    def generate_vql_triage_pack(self, target_os="windows", output_file=None):
        """Generate VQL artifact collection pack for triage."""
        artifacts = (
            self.TRIAGE_ARTIFACTS_WINDOWS
            if target_os == "windows"
            else self.TRIAGE_ARTIFACTS_LINUX
        )

        vql_queries = []
        for artifact in artifacts:
            vql_queries.append({
                "artifact": artifact,
                "vql": f"SELECT * FROM Artifact.{artifact}()",
                "description": f"Collect {artifact.split('.')[-1]} artifacts",
            })

        pack = {
            "name": f"IR Triage Pack - {target_os.title()}",
            "version": "1.0",
            "generated": datetime.utcnow().isoformat(),
            "target_os": target_os,
            "artifacts": vql_queries,
        }

        if output_file is None:
            output_file = self.output_dir / f"triage_pack_{target_os}.json"

        with open(output_file, "w") as f:
            json.dump(pack, f, indent=2)

        print(f"[+] Generated {target_os} triage pack: {output_file}")
        print(f"    Artifacts: {len(vql_queries)}")
        return pack

    def generate_hunt_config(self, hunt_name, description, artifacts, parameters=None,
                              include_labels=None, exclude_labels=None):
        """Generate hunt configuration for enterprise-wide collection."""
        hunt_config = {
            "hunt_name": hunt_name,
            "description": description,
            "created": datetime.utcnow().isoformat(),
            "artifacts": artifacts,
            "parameters": parameters or {},
            "targeting": {
                "include_labels": include_labels or [],
                "exclude_labels": exclude_labels or [],
            },
            "resource_limits": {
                "cpu_limit": 50,
                "iops_limit": 100,
                "timeout_seconds": 600,
                "max_rows": 100000,
                "max_upload_bytes": 104857600,
            },
        }

        config_file = self.output_dir / f"hunt_{hunt_name.replace(' ', '_')}.json"
        with open(config_file, "w") as f:
            json.dump(hunt_config, f, indent=2)

        print(f"[+] Generated hunt config: {config_file}")
        return hunt_config

    def generate_ioc_hunt(self, iocs, hunt_name="IOC Hunt"):
        """Generate VQL-based IOC hunt queries."""
        vql_queries = []

        if "hashes" in iocs:
            hash_list = "|".join(iocs["hashes"])
            vql_queries.append({
                "name": "Hash Hunt",
                "vql": f"""SELECT * FROM Artifact.Generic.Detection.HashHunter(
                    Hashes='{hash_list}'
                )""",
            })

        if "ips" in iocs:
            ip_regex = "|".join(ip.replace(".", "\\.") for ip in iocs["ips"])
            vql_queries.append({
                "name": "Network Connection Hunt",
                "vql": f"""SELECT * FROM Artifact.Windows.Network.Netstat()
                    WHERE RemoteAddr =~ '{ip_regex}'""",
            })

        if "domains" in iocs:
            domain_regex = "|".join(d.replace(".", "\\.") for d in iocs["domains"])
            vql_queries.append({
                "name": "DNS Cache Hunt",
                "vql": f"""SELECT * FROM Artifact.Windows.Network.DNSCache()
                    WHERE Name =~ '{domain_regex}'""",
            })

        if "filenames" in iocs:
            file_regex = "|".join(iocs["filenames"])
            vql_queries.append({
                "name": "File Hunt",
                "vql": f"""SELECT * FROM Artifact.Windows.NTFS.MFT(
                    FileRegex='{file_regex}'
                )""",
            })

        if "yara_rules" in iocs:
            for rule_name, rule_content in iocs["yara_rules"].items():
                vql_queries.append({
                    "name": f"YARA Hunt - {rule_name}",
                    "vql": f"""SELECT * FROM Artifact.Windows.Detection.Yara.Process(
                        YaraRule='{rule_content}'
                    )""",
                })

        hunt_file = self.output_dir / f"ioc_hunt_{hunt_name.replace(' ', '_')}.json"
        with open(hunt_file, "w") as f:
            json.dump({"name": hunt_name, "queries": vql_queries}, f, indent=2)

        print(f"[+] Generated IOC hunt with {len(vql_queries)} queries: {hunt_file}")
        return vql_queries

    def generate_collection_checklist(self, case_id, target_hosts):
        """Generate collection checklist for IR case."""
        checklist = {
            "case_id": case_id,
            "generated": datetime.utcnow().isoformat(),
            "targets": [],
        }

        for host in target_hosts:
            target = {
                "hostname": host,
                "collections": [
                    {"artifact": "Volatile Data", "status": "pending", "items": [
                        "Running processes", "Network connections", "DNS cache",
                        "Logged-in users", "Open files",
                    ]},
                    {"artifact": "Event Logs", "status": "pending", "items": [
                        "Security.evtx", "System.evtx", "Application.evtx",
                        "PowerShell/Operational.evtx", "Sysmon/Operational.evtx",
                    ]},
                    {"artifact": "Execution Evidence", "status": "pending", "items": [
                        "Prefetch files", "Amcache.hve", "Shimcache",
                        "UserAssist", "BAM/DAM",
                    ]},
                    {"artifact": "Persistence Mechanisms", "status": "pending", "items": [
                        "Scheduled tasks", "Services", "Registry Run keys",
                        "WMI subscriptions", "Startup items",
                    ]},
                    {"artifact": "File System", "status": "pending", "items": [
                        "MFT entries", "USN Journal", "Recycle Bin",
                        "Recent files", "Downloads folder",
                    ]},
                    {"artifact": "User Activity", "status": "pending", "items": [
                        "Browser history", "PowerShell history", "RDP cache",
                        "Recent documents", "Jump lists",
                    ]},
                ],
            }
            checklist["targets"].append(target)

        checklist_file = self.output_dir / f"collection_checklist_{case_id}.json"
        with open(checklist_file, "w") as f:
            json.dump(checklist, f, indent=2)

        print(f"[+] Collection checklist for {len(target_hosts)} hosts: {checklist_file}")
        return checklist

    def analyze_collection_results(self, results_dir):
        """Analyze collected Velociraptor results for suspicious indicators."""
        results_path = Path(results_dir)
        findings = []

        for json_file in results_path.glob("**/*.json"):
            try:
                with open(json_file) as f:
                    data = json.load(f)
                if isinstance(data, list):
                    for item in data:
                        finding = self._check_for_indicators(item, str(json_file))
                        if finding:
                            findings.append(finding)
            except (json.JSONDecodeError, KeyError):
                continue

        findings_file = self.output_dir / "analysis_findings.json"
        with open(findings_file, "w") as f:
            json.dump(findings, f, indent=2, default=str)

        print(f"[+] Analysis complete: {len(findings)} findings -> {findings_file}")
        return findings

    def _check_for_indicators(self, item, source_file):
        """Check a single result item for suspicious indicators."""
        suspicious_processes = [
            "mimikatz", "rubeus", "lazagne", "sharphound", "bloodhound",
            "cobaltstrike", "beacon", "psexec", "wmiexec", "smbexec",
        ]
        suspicious_commands = [
            "invoke-mimikatz", "invoke-expression", "downloadstring",
            "net user /add", "net localgroup administrators",
            "reg save hklm\\sam", "reg save hklm\\system",
            "ntdsutil", "vssadmin create shadow",
        ]

        name = str(item.get("Name", "") or item.get("CommandLine", "")).lower()

        for proc in suspicious_processes:
            if proc in name:
                return {
                    "severity": "CRITICAL",
                    "indicator": f"Suspicious process: {proc}",
                    "detail": item,
                    "source": source_file,
                }

        for cmd in suspicious_commands:
            if cmd in name:
                return {
                    "severity": "HIGH",
                    "indicator": f"Suspicious command: {cmd}",
                    "detail": item,
                    "source": source_file,
                }

        return None

    def generate_report(self):
        """Generate collection summary report."""
        report = {
            "title": "Velociraptor IR Collection Report",
            "generated": datetime.utcnow().isoformat(),
            "collections": self.collection_manifest,
            "output_directory": str(self.output_dir),
        }

        report_file = self.output_dir / "collection_report.json"
        with open(report_file, "w") as f:
            json.dump(report, f, indent=2)

        print(f"[+] Collection report: {report_file}")
        return report


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Velociraptor IR Collection Manager")
    parser.add_argument("--action", choices=[
        "triage-pack", "hunt-config", "ioc-hunt", "checklist", "analyze",
    ], required=True)
    parser.add_argument("--os", default="windows", choices=["windows", "linux"])
    parser.add_argument("--case-id", default="IR-2025-001")
    parser.add_argument("--hosts", nargs="+", help="Target hostnames")
    parser.add_argument("--iocs", help="JSON file with IOCs")
    parser.add_argument("--results-dir", help="Directory with collection results")
    parser.add_argument("-o", "--output", default="velociraptor_results")

    args = parser.parse_args()
    collector = VelociraptorCollector(output_dir=args.output)

    if args.action == "triage-pack":
        collector.generate_vql_triage_pack(target_os=args.os)
    elif args.action == "hunt-config":
        collector.generate_hunt_config("IR Hunt", "Incident response hunt",
                                        collector.TRIAGE_ARTIFACTS_WINDOWS)
    elif args.action == "ioc-hunt":
        if args.iocs:
            with open(args.iocs) as f:
                iocs = json.load(f)
            collector.generate_ioc_hunt(iocs)
    elif args.action == "checklist":
        hosts = args.hosts or ["WKS001", "SRV001", "DC01"]
        collector.generate_collection_checklist(args.case_id, hosts)
    elif args.action == "analyze":
        if args.results_dir:
            collector.analyze_collection_results(args.results_dir)


if __name__ == "__main__":
    main()