sica-fondt/knowledge/cyber-skills/skills/building-automated-malware-submission-pipeline/scripts/agent.py

#!/usr/bin/env python3
"""Agent for building an automated malware submission and analysis pipeline."""

import os
import json
import hashlib
import time
import argparse
from datetime import datetime
from pathlib import Path

import requests


def compute_hashes(filepath):
    """Calculate MD5, SHA1, SHA256 hashes and file size."""
    with open(filepath, "rb") as f:
        content = f.read()
    return {
        "md5": hashlib.md5(content).hexdigest(),
        "sha1": hashlib.sha1(content).hexdigest(),
        "sha256": hashlib.sha256(content).hexdigest(),
        "size": len(content),
    }


def check_virustotal(sha256, api_key):
    """Check if a file hash is known in VirusTotal."""
    resp = requests.get(
        f"https://www.virustotal.com/api/v3/files/{sha256}",
        headers={"x-apikey": api_key},
        timeout=30,
    )
    if resp.status_code == 200:
        attrs = resp.json()["data"]["attributes"]
        stats = attrs.get("last_analysis_stats", {})
        ptc = attrs.get("popular_threat_classification", {})
        return {
            "found": True,
            "malicious": stats.get("malicious", 0),
            "suspicious": stats.get("suspicious", 0),
            "total": sum(stats.values()),
            "threat_label": ptc.get("suggested_threat_label", "Unknown"),
            "type_description": attrs.get("type_description", "Unknown"),
        }
    return {"found": False}


def check_malwarebazaar(sha256):
    """Check if a hash is known in MalwareBazaar."""
    resp = requests.post(
        "https://mb-api.abuse.ch/api/v1/",
        data={"query": "get_info", "hash": sha256},
        timeout=30,
    )
    if resp.status_code == 200:
        data = resp.json()
        if data.get("query_status") == "ok" and data.get("data"):
            entry = data["data"][0]
            return {
                "found": True,
                "signature": entry.get("signature", "Unknown"),
                "tags": entry.get("tags", []),
                "file_type": entry.get("file_type", "Unknown"),
                "first_seen": entry.get("first_seen", "Unknown"),
            }
    return {"found": False}


def submit_to_cuckoo(filepath, cuckoo_url, timeout=300):
    """Submit a file to Cuckoo Sandbox for dynamic analysis."""
    with open(filepath, "rb") as f:
        resp = requests.post(
            f"{cuckoo_url}/tasks/create/file",
            files={"file": f},
            data={"timeout": timeout, "options": "procmemdump=yes", "priority": 2},
        )
    if resp.status_code == 200:
        return resp.json().get("task_id")
    return None


def wait_for_cuckoo_report(task_id, cuckoo_url, poll_interval=30, max_wait=600):
    """Poll Cuckoo until analysis is complete, then return the report."""
    elapsed = 0
    while elapsed < max_wait:
        resp = requests.get(f"{cuckoo_url}/tasks/view/{task_id}", timeout=30)
        if resp.status_code == 200:
            status = resp.json().get("task", {}).get("status")
            if status == "reported":
                report_resp = requests.get(f"{cuckoo_url}/tasks/report/{task_id}", timeout=30)
                return report_resp.json() if report_resp.status_code == 200 else None
            if status == "failed_analysis":
                return None
        time.sleep(poll_interval)
        elapsed += poll_interval
    return None


def extract_iocs_from_report(report):
    """Extract IOCs from a Cuckoo sandbox report."""
    iocs = {"ips": [], "domains": [], "urls": [], "hashes": [], "registry_keys": []}
    if not report:
        return iocs
    network = report.get("network", {})
    iocs["domains"] = [d.get("request", "") for d in network.get("dns", [])]
    iocs["ips"] = network.get("hosts", [])
    iocs["urls"] = [h.get("uri", "") for h in network.get("http", [])]
    iocs["hashes"] = [f.get("sha256", "") for f in report.get("dropped", []) if f.get("sha256")]
    behavior = report.get("behavior", {}).get("summary", {})
    iocs["registry_keys"] = behavior.get("regkey_written", [])[:10]
    return iocs


def generate_verdict(vt_result, mb_result, sandbox_report):
    """Generate a verdict combining pre-screening and sandbox results."""
    vt_malicious = vt_result.get("malicious", 0) if vt_result.get("found") else 0
    sandbox_score = sandbox_report.get("info", {}).get("score", 0) if sandbox_report else 0
    sig_count = len(sandbox_report.get("signatures", [])) if sandbox_report else 0
    combined = (vt_malicious * 2) + (sandbox_score * 10) + (sig_count * 5)
    if combined >= 100:
        return "MALICIOUS", "HIGH", combined
    elif combined >= 50:
        return "SUSPICIOUS", "MEDIUM", combined
    elif combined >= 20:
        return "POTENTIALLY_UNWANTED", "LOW", combined
    return "CLEAN", "HIGH", combined


def push_to_splunk_hec(event_data, splunk_url, splunk_token):
    """Send analysis results to Splunk via HTTP Event Collector."""
    payload = {
        "sourcetype": "malware_analysis",
        "source": "malware_pipeline",
        "event": event_data,
    }
    resp = requests.post(
        f"{splunk_url}/services/collector/event",
        headers={"Authorization": f"Splunk {splunk_token}"},
        json=payload,
        verify=not os.environ.get("SKIP_TLS_VERIFY", "").lower() == "true",  # Set SKIP_TLS_VERIFY=true for self-signed certs in lab environments
        timeout=30,
    )
    return resp.status_code == 200


def main():
    parser = argparse.ArgumentParser(description="Malware Submission Pipeline Agent")
    parser.add_argument("--sample", required=True, help="Path to sample file")
    parser.add_argument("--vt-key", default=os.getenv("VT_API_KEY"))
    parser.add_argument("--cuckoo-url", default=os.getenv("CUCKOO_URL", "http://localhost:8090"))
    parser.add_argument("--splunk-url", default=os.getenv("SPLUNK_HEC_URL"))
    parser.add_argument("--splunk-token", default=os.getenv("SPLUNK_HEC_TOKEN"))
    parser.add_argument("--output", default="malware_report.json")
    parser.add_argument("--skip-sandbox", action="store_true")
    args = parser.parse_args()

    print(f"[+] Analyzing sample: {args.sample}")
    hashes = compute_hashes(args.sample)
    print(f"    SHA256: {hashes['sha256']}")

    print("[+] Pre-screening with VirusTotal...")
    vt_result = check_virustotal(hashes["sha256"], args.vt_key) if args.vt_key else {"found": False}
    if vt_result.get("found"):
        print(f"    VT: {vt_result['malicious']}/{vt_result['total']} detections")

    print("[+] Checking MalwareBazaar...")
    mb_result = check_malwarebazaar(hashes["sha256"])
    if mb_result.get("found"):
        print(f"    MB: {mb_result['signature']} ({mb_result['file_type']})")

    sandbox_report = None
    if not args.skip_sandbox and (not vt_result.get("found") or vt_result.get("malicious", 0) < 5):
        print("[+] Submitting to Cuckoo Sandbox...")
        task_id = submit_to_cuckoo(args.sample, args.cuckoo_url)
        if task_id:
            print(f"    Task ID: {task_id}, waiting for analysis...")
            sandbox_report = wait_for_cuckoo_report(task_id, args.cuckoo_url)

    verdict, confidence, score = generate_verdict(vt_result, mb_result, sandbox_report)
    iocs = extract_iocs_from_report(sandbox_report)

    report = {
        "sha256": hashes["sha256"],
        "hashes": hashes,
        "virustotal": vt_result,
        "malwarebazaar": mb_result,
        "verdict": verdict,
        "confidence": confidence,
        "combined_score": score,
        "iocs": iocs,
        "analyzed_at": datetime.utcnow().isoformat(),
    }
    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    print(f"\n[+] VERDICT: {verdict} (confidence: {confidence}, score: {score})")
    print(f"[+] Report saved to {args.output}")

    if args.splunk_url and args.splunk_token:
        if push_to_splunk_hec(report, args.splunk_url, args.splunk_token):
            print("[+] Results pushed to Splunk")


if __name__ == "__main__":
    main()