#!/usr/bin/env python3 """Agent for automating IOC enrichment with VirusTotal, AbuseIPDB, and STIX.""" import os import re import json import time import argparse from datetime import datetime from dataclasses import dataclass, field import requests from stix2 import Indicator, Bundle RATE_LIMIT_DELAY = 0.25 @dataclass class EnrichmentResult: ioc_value: str ioc_type: str vt_malicious: int = 0 vt_total: int = 0 vt_threat_label: str = "" abuse_confidence: int = 0 abuse_reports: int = 0 shodan_ports: list = field(default_factory=list) confidence_score: int = 0 def classify_ioc(value): """Auto-detect IOC type from value.""" if re.match(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$", value): return "ip" if re.match(r"^[a-fA-F0-9]{64}$", value): return "sha256" if re.match(r"^[a-fA-F0-9]{32}$", value): return "md5" if re.match(r"^https?://", value): return "url" return "domain" def enrich_ip_virustotal(ip, api_key): """Enrich an IP address via VirusTotal API v3.""" resp = requests.get( f"https://www.virustotal.com/api/v3/ip_addresses/{ip}", headers={"x-apikey": api_key}, timeout=30, ) if resp.status_code == 200: attrs = resp.json()["data"]["attributes"] stats = attrs.get("last_analysis_stats", {}) return { "malicious": stats.get("malicious", 0), "total": sum(stats.values()), "country": attrs.get("country", ""), "asn": attrs.get("asn", 0), "as_owner": attrs.get("as_owner", ""), } return {} def enrich_hash_virustotal(file_hash, api_key): """Enrich a file hash via VirusTotal API v3.""" resp = requests.get( f"https://www.virustotal.com/api/v3/files/{file_hash}", 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 { "malicious": stats.get("malicious", 0), "total": sum(stats.values()), "threat_label": ptc.get("suggested_threat_label", ""), "type_description": attrs.get("type_description", ""), } return {} def enrich_domain_virustotal(domain, api_key): """Enrich a domain via VirusTotal API v3.""" resp = requests.get( f"https://www.virustotal.com/api/v3/domains/{domain}", headers={"x-apikey": api_key}, timeout=30, ) if resp.status_code == 200: attrs = resp.json()["data"]["attributes"] stats = attrs.get("last_analysis_stats", {}) return { "malicious": stats.get("malicious", 0), "total": sum(stats.values()), "registrar": attrs.get("registrar", ""), } return {} def enrich_ip_abuseipdb(ip, api_key): """Check an IP against AbuseIPDB.""" resp = requests.get( "https://api.abuseipdb.com/api/v2/check", headers={"Key": api_key, "Accept": "application/json"}, params={"ipAddress": ip, "maxAgeInDays": 90}, timeout=30, ) if resp.status_code == 200: data = resp.json()["data"] return { "abuse_confidence": data.get("abuseConfidenceScore", 0), "total_reports": data.get("totalReports", 0), "country": data.get("countryCode", ""), "isp": data.get("isp", ""), } return {} def compute_confidence(vt_result, abuse_result=None): """Calculate composite confidence score from enrichment sources.""" vt_score = 0 if vt_result.get("total", 0) > 0: vt_score = (vt_result["malicious"] / vt_result["total"]) * 60 abuse_score = 0 if abuse_result: abuse_score = (abuse_result.get("abuse_confidence", 0) / 100) * 40 return min(int(vt_score + abuse_score), 100) def enrich_ioc(value, ioc_type, vt_key, abuse_key=None): """Enrich a single IOC through all available sources.""" result = EnrichmentResult(ioc_value=value, ioc_type=ioc_type) vt_data = {} if ioc_type == "ip": vt_data = enrich_ip_virustotal(value, vt_key) time.sleep(RATE_LIMIT_DELAY) if abuse_key: abuse_data = enrich_ip_abuseipdb(value, abuse_key) result.abuse_confidence = abuse_data.get("abuse_confidence", 0) result.abuse_reports = abuse_data.get("total_reports", 0) elif ioc_type in ("sha256", "md5"): vt_data = enrich_hash_virustotal(value, vt_key) time.sleep(RATE_LIMIT_DELAY) elif ioc_type == "domain": vt_data = enrich_domain_virustotal(value, vt_key) time.sleep(RATE_LIMIT_DELAY) result.vt_malicious = vt_data.get("malicious", 0) result.vt_total = vt_data.get("total", 0) result.vt_threat_label = vt_data.get("threat_label", "") abuse_dict = {"abuse_confidence": result.abuse_confidence} if ioc_type == "ip" else None result.confidence_score = compute_confidence(vt_data, abuse_dict) return result def export_stix_indicators(results, output_path): """Export enriched IOCs as STIX 2.1 indicators.""" pattern_map = { "ip": lambda v: f"[ipv4-addr:value = '{v}']", "domain": lambda v: f"[domain-name:value = '{v}']", "sha256": lambda v: f"[file:hashes.'SHA-256' = '{v}']", "md5": lambda v: f"[file:hashes.MD5 = '{v}']", "url": lambda v: f"[url:value = '{v}']", } indicators = [] for r in results: pattern_fn = pattern_map.get(r.ioc_type) if pattern_fn: ind = Indicator( name=f"{r.ioc_type}: {r.ioc_value}", pattern=pattern_fn(r.ioc_value), pattern_type="stix", valid_from=datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"), confidence=r.confidence_score, ) indicators.append(ind) bundle = Bundle(objects=indicators, allow_custom=True) with open(output_path, "w") as f: f.write(bundle.serialize(pretty=True)) return len(indicators) def main(): parser = argparse.ArgumentParser(description="IOC Enrichment Automation Agent") parser.add_argument("--vt-key", default=os.getenv("VT_API_KEY"), help="VirusTotal API key") parser.add_argument("--abuse-key", default=os.getenv("ABUSEIPDB_KEY"), help="AbuseIPDB API key") parser.add_argument("--ioc-file", help="File with IOCs (one per line)") parser.add_argument("--ioc", help="Single IOC to enrich") parser.add_argument("--output", default="enrichment_results.json") parser.add_argument("--stix-output", help="Export as STIX bundle") args = parser.parse_args() iocs = [] if args.ioc: iocs.append(args.ioc) if args.ioc_file: with open(args.ioc_file) as f: iocs.extend(line.strip() for line in f if line.strip() and not line.startswith("#")) results = [] for ioc_val in iocs: ioc_type = classify_ioc(ioc_val) print(f" Enriching {ioc_type}: {ioc_val}...") result = enrich_ioc(ioc_val, ioc_type, args.vt_key, args.abuse_key) results.append(result) verdict = "MALICIOUS" if result.confidence_score >= 70 else "SUSPICIOUS" if result.confidence_score >= 40 else "CLEAN" print(f" VT: {result.vt_malicious}/{result.vt_total} | Confidence: {result.confidence_score} | {verdict}") report = { "enriched_at": datetime.utcnow().isoformat(), "total_iocs": len(results), "results": [r.__dict__ for r in results], } with open(args.output, "w") as f: json.dump(report, f, indent=2, default=str) print(f"[+] Results saved to {args.output}") if args.stix_output: count = export_stix_indicators(results, args.stix_output) print(f"[+] Exported {count} STIX indicators to {args.stix_output}") if __name__ == "__main__": main()