#!/usr/bin/env python3 """ malware_classifier.py - Classify malware using feature extraction and similarity scoring. Extracts static and behavioral features from binaries, computes similarity metrics, clusters related samples, and maintains a classification database. Usage: python3 malware_classifier.py --extract-features --input sample.exe --output features.json python3 malware_classifier.py --compare --input a.exe --reference b.exe --output cmp.json python3 malware_classifier.py --cluster --input features.json --algorithm dbscan --output clusters.json python3 malware_classifier.py --classify --input sample.exe --database db.json --output result.json python3 malware_classifier.py --add-to-db --input sample.exe --family emotet --database db.json """ from __future__ import annotations import argparse import hashlib import json import logging import math import os import struct import sys from collections import Counter from datetime import datetime, timezone from pathlib import Path from typing import Optional logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- # PE signature constants MZ_SIGNATURE = b"MZ" PE_SIGNATURE = b"PE\x00\x00" # Suspicious API imports commonly found in malware SUSPICIOUS_APIS = { "CreateRemoteThread", "VirtualAllocEx", "WriteProcessMemory", "NtUnmapViewOfSection", "QueueUserAPC", "SetWindowsHookEx", "CreateToolhelp32Snapshot", "Process32First", "Process32Next", "OpenProcess", "IsDebuggerPresent", "NtQueryInformationProcess", "GetTickCount", "QueryPerformanceCounter", "CryptEncrypt", "CryptDecrypt", "InternetOpen", "HttpSendRequest", "URLDownloadToFile", "WinExec", "ShellExecute", "CreateService", "RegSetValueEx", "NtCreateThreadEx", "RtlCreateUserThread", } # Common malware-associated DLLs NETWORK_DLLS = {"ws2_32.dll", "wininet.dll", "winhttp.dll", "urlmon.dll"} CRYPTO_DLLS = {"advapi32.dll", "bcrypt.dll", "ncrypt.dll", "crypt32.dll"} # --------------------------------------------------------------------------- # PE parsing helpers # --------------------------------------------------------------------------- def compute_file_hashes(file_path: str) -> dict: """Compute MD5, SHA1, and SHA256 hashes of a file. Args: file_path: Path to the file. Returns: Dictionary with hash algorithm names as keys and hex digests as values. """ hashers = { "md5": hashlib.md5(), "sha1": hashlib.sha1(), "sha256": hashlib.sha256(), } with open(file_path, "rb") as f: while True: chunk = f.read(65536) if not chunk: break for h in hashers.values(): h.update(chunk) return {name: h.hexdigest() for name, h in hashers.items()} def compute_entropy(data: bytes) -> float: """Compute Shannon entropy of a byte sequence. Args: data: Byte sequence to analyze. Returns: Entropy value between 0.0 (uniform) and 8.0 (maximum randomness). """ if not data: return 0.0 counts = Counter(data) length = len(data) entropy = 0.0 for count in counts.values(): if count == 0: continue p = count / length entropy -= p * math.log2(p) return round(entropy, 4) def compute_imphash(imports: dict) -> str: """Compute import hash (imphash) from import table data. The imphash is an MD5 of the ordered, lowercased DLL.function pairs from the import table. Samples built from the same source code with the same imports produce identical imphashes. Args: imports: Dictionary mapping DLL names to lists of function names. Returns: Hex digest string of the import hash. """ entries = [] for dll_name in sorted(imports.keys()): dll_base = dll_name.lower().replace(".dll", "") for func in imports[dll_name]: entries.append(f"{dll_base}.{func.lower()}") combined = ",".join(entries) return hashlib.md5(combined.encode()).hexdigest() def parse_pe_header(data: bytes) -> dict: """Parse PE header fields from raw binary data. Extracts machine type, timestamp, section count, characteristics, entry point, and subsystem from the PE COFF and optional headers. Args: data: Raw bytes of the PE file (at least first 4KB). Returns: Dictionary with parsed PE header fields. """ result = {} if data[:2] != MZ_SIGNATURE: return {"error": "Not a PE file (missing MZ signature)"} # PE header offset from DOS header at 0x3C pe_offset = struct.unpack_from(" list: """Parse PE section headers. Extracts section names, sizes, virtual addresses, and computes entropy for each section's raw data. Args: data: Raw bytes of the PE file. Returns: List of dictionaries with section properties. """ if data[:2] != MZ_SIGNATURE: return [] pe_offset = struct.unpack_from(" len(data): break name_bytes = data[sec_offset:sec_offset + 8] name = name_bytes.split(b"\x00")[0].decode("ascii", errors="replace") virtual_size = struct.unpack_from(" 0 else b"" entropy = compute_entropy(section_data) sec_hash = hashlib.sha256(section_data).hexdigest() if section_data else "" sections.append({ "name": name, "virtual_size": virtual_size, "virtual_address": hex(virtual_addr), "raw_size": raw_size, "entropy": entropy, "sha256": sec_hash, "characteristics": hex(characteristics), "executable": bool(characteristics & 0x20000000), "writable": bool(characteristics & 0x80000000), }) return sections def extract_strings(data: bytes, min_length: int = 6) -> list: """Extract printable ASCII strings from binary data. Args: data: Raw binary data. min_length: Minimum string length to extract. Returns: List of extracted strings. """ strings = [] current = [] for byte in data: if 32 <= byte < 127: current.append(chr(byte)) else: if len(current) >= min_length: strings.append("".join(current)) current = [] if len(current) >= min_length: strings.append("".join(current)) return strings def extract_import_table(data: bytes) -> dict: """Extract import table from PE binary. This is a simplified parser that searches for DLL name patterns in the binary. For full IAT parsing, use pefile library. Args: data: Raw PE file bytes. Returns: Dictionary mapping DLL names to lists of imported function names. """ # Simplified approach: extract DLL-like and API-like strings strings = extract_strings(data, min_length=4) imports = {} current_dll = None for s in strings: s_lower = s.lower() if s_lower.endswith(".dll"): current_dll = s if current_dll not in imports: imports[current_dll] = [] elif current_dll and s[0].isupper() and len(s) > 3 and not s.startswith("\\"): # Heuristic: API names are typically CamelCase imports[current_dll].append(s) return imports # --------------------------------------------------------------------------- # Feature extraction # --------------------------------------------------------------------------- def extract_static_features(file_path: str) -> dict: """Extract static analysis features from a PE binary. Parses PE headers, sections, imports, strings, and computes various hashes and metrics used for classification. Args: file_path: Path to the PE binary. Returns: Dictionary containing all extracted static features. """ path = Path(file_path) data = path.read_bytes() hashes = compute_file_hashes(file_path) header = parse_pe_header(data) sections = parse_pe_sections(data) imports = extract_import_table(data) strings = extract_strings(data) # Compute imphash imphash = compute_imphash(imports) if imports else "" # Rich header hash (simplified: hash bytes between "DanS" and PE signature) rich_hash = "" dans_offset = data.find(b"DanS") if dans_offset > 0: pe_offset = struct.unpack_from(" 5] registry_keys = [s for s in strings if "SOFTWARE\\" in s or "CurrentVersion\\" in s] return { "file_name": path.name, "file_size": len(data), "file_entropy": file_entropy, "hashes": hashes, "imphash": imphash, "rich_header_hash": rich_hash, "pe_header": header, "sections": sections, "section_names": [s["name"] for s in sections], "section_entropies": [s["entropy"] for s in sections], "import_count": sum(len(v) for v in imports.values()), "import_dlls": sorted(imports.keys()), "suspicious_imports": suspicious, "has_network_imports": has_network, "has_crypto_imports": has_crypto, "string_count": len(strings), "url_strings": urls[:50], "file_path_strings": file_paths[:50], "registry_strings": registry_keys[:50], "extracted_at": datetime.now(timezone.utc).isoformat(), } def extract_behavioral_features(report_path: str, fmt: str = "json") -> dict: """Extract behavioral features from a sandbox report or API trace. Parses dynamic analysis output to extract API call sequences, network behavior, file operations, and other behavioral indicators. Args: report_path: Path to sandbox report (JSON) or API trace log. fmt: Input format - 'json' for sandbox reports, 'strace' for traces. Returns: Dictionary containing behavioral feature vectors. """ path = Path(report_path) if not path.exists(): raise FileNotFoundError(f"Report not found: {report_path}") if fmt == "json": with open(path) as f: report = json.load(f) # Generic extraction from common sandbox report formats behavior = report.get("behavior", report.get("dynamic", {})) processes = behavior.get("processes", behavior.get("process_list", [])) network = report.get("network", behavior.get("network", {})) api_calls = [] for proc in processes: calls = proc.get("calls", proc.get("api_calls", [])) for call in calls: api_name = call.get("api", call.get("name", "")) if api_name: api_calls.append(api_name) dns_queries = [d.get("request", d.get("hostname", "")) for d in network.get("dns", [])] http_requests = [h.get("uri", h.get("url", "")) for h in network.get("http", [])] return { "api_call_count": len(api_calls), "unique_apis": len(set(api_calls)), "api_frequency": dict(Counter(api_calls).most_common(30)), "network_connections": len(network.get("tcp", []) + network.get("udp", [])), "dns_queries": dns_queries[:50], "http_requests": http_requests[:50], "processes_spawned": len(processes), "files_created": len(behavior.get("files_created", [])), "files_modified": len(behavior.get("files_modified", [])), "registry_keys_modified": len(behavior.get("registry_modified", [])), "extracted_at": datetime.now(timezone.utc).isoformat(), } elif fmt == "strace": text = path.read_text() syscalls = [] for line in text.splitlines(): if "(" in line: call_name = line.split("(")[0].strip() if call_name and not call_name.startswith("#"): syscalls.append(call_name) return { "syscall_count": len(syscalls), "unique_syscalls": len(set(syscalls)), "syscall_frequency": dict(Counter(syscalls).most_common(30)), "extracted_at": datetime.now(timezone.utc).isoformat(), } return {"error": f"Unsupported format: {fmt}"} # --------------------------------------------------------------------------- # Similarity and comparison # --------------------------------------------------------------------------- def jaccard_similarity(set_a: set, set_b: set) -> float: """Compute Jaccard similarity between two sets. Args: set_a: First set. set_b: Second set. Returns: Jaccard index between 0.0 and 1.0. """ if not set_a and not set_b: return 1.0 intersection = len(set_a & set_b) union = len(set_a | set_b) return round(intersection / union, 4) if union > 0 else 0.0 def cosine_similarity(vec_a: list, vec_b: list) -> float: """Compute cosine similarity between two numeric vectors. Args: vec_a: First feature vector. vec_b: Second feature vector. Returns: Cosine similarity between -1.0 and 1.0. """ if len(vec_a) != len(vec_b): min_len = min(len(vec_a), len(vec_b)) vec_a = vec_a[:min_len] vec_b = vec_b[:min_len] dot_product = sum(a * b for a, b in zip(vec_a, vec_b)) mag_a = math.sqrt(sum(a * a for a in vec_a)) mag_b = math.sqrt(sum(b * b for b in vec_b)) if mag_a == 0 or mag_b == 0: return 0.0 return round(dot_product / (mag_a * mag_b), 4) def compare_samples(features_a: dict, features_b: dict) -> dict: """Compare two samples using multiple similarity metrics. Args: features_a: Feature dictionary for sample A. features_b: Feature dictionary for sample B. Returns: Dictionary with per-metric similarity scores and overall score. """ scores = {} # Imphash match imphash_a = features_a.get("imphash", "") imphash_b = features_b.get("imphash", "") scores["imphash_match"] = 1.0 if (imphash_a and imphash_a == imphash_b) else 0.0 # Rich header hash match rich_a = features_a.get("rich_header_hash", "") rich_b = features_b.get("rich_header_hash", "") scores["rich_header_match"] = 1.0 if (rich_a and rich_a == rich_b) else 0.0 # Import set similarity imports_a = set(features_a.get("import_dlls", [])) imports_b = set(features_b.get("import_dlls", [])) scores["import_dll_jaccard"] = jaccard_similarity(imports_a, imports_b) # Section name similarity sections_a = set(features_a.get("section_names", [])) sections_b = set(features_b.get("section_names", [])) scores["section_name_jaccard"] = jaccard_similarity(sections_a, sections_b) # Section entropy vector similarity ent_a = features_a.get("section_entropies", []) ent_b = features_b.get("section_entropies", []) if ent_a and ent_b: scores["entropy_cosine"] = cosine_similarity(ent_a, ent_b) else: scores["entropy_cosine"] = 0.0 # Suspicious import overlap susp_a = set(features_a.get("suspicious_imports", [])) susp_b = set(features_b.get("suspicious_imports", [])) scores["suspicious_import_jaccard"] = jaccard_similarity(susp_a, susp_b) # Weighted overall score weights = { "imphash_match": 0.30, "rich_header_match": 0.10, "import_dll_jaccard": 0.20, "section_name_jaccard": 0.10, "entropy_cosine": 0.15, "suspicious_import_jaccard": 0.15, } overall = sum(scores[k] * weights[k] for k in weights) scores["overall_similarity"] = round(overall, 4) return scores # --------------------------------------------------------------------------- # Classification database # --------------------------------------------------------------------------- def init_database(output_path: str) -> dict: """Initialize a new classification database. Args: output_path: Path to write the database JSON. Returns: Empty database dictionary. """ db = { "created": datetime.now(timezone.utc).isoformat(), "version": "1.0", "samples": [], "families": {}, "total_samples": 0, } Path(output_path).write_text(json.dumps(db, indent=2)) logger.info(f"Database initialized: {output_path}") return db def add_to_database( database_path: str, features: dict, family: str, malware_type: str = "unknown", campaign: str = "", confidence: float = 1.0, ) -> dict: """Add a classified sample to the database. Args: database_path: Path to the database JSON file. features: Extracted feature dictionary for the sample. family: Malware family name. malware_type: Malware type (trojan, ransomware, etc.). campaign: Campaign identifier. confidence: Classification confidence score (0.0-1.0). Returns: The database entry that was added. """ db_path = Path(database_path) db = json.loads(db_path.read_text()) entry = { "sha256": features.get("hashes", {}).get("sha256", ""), "file_name": features.get("file_name", ""), "family": family, "type": malware_type, "campaign": campaign, "confidence": confidence, "imphash": features.get("imphash", ""), "rich_header_hash": features.get("rich_header_hash", ""), "import_dlls": features.get("import_dlls", []), "section_names": features.get("section_names", []), "section_entropies": features.get("section_entropies", []), "suspicious_imports": features.get("suspicious_imports", []), "added_at": datetime.now(timezone.utc).isoformat(), } db["samples"].append(entry) db["total_samples"] = len(db["samples"]) # Update family counts if family not in db["families"]: db["families"][family] = 0 db["families"][family] += 1 db_path.write_text(json.dumps(db, indent=2)) logger.info(f"Added {entry['sha256'][:16]}... as {family} to database") return entry def classify_against_database( features: dict, database_path: str, top_k: int = 5, threshold: float = 0.5, ) -> dict: """Classify a sample by comparing against the database. Args: features: Extracted features of the sample to classify. database_path: Path to the classification database. top_k: Number of top matches to return. threshold: Minimum similarity score to consider a match. Returns: Classification result with top matches and confidence. """ db = json.loads(Path(database_path).read_text()) matches = [] for db_sample in db.get("samples", []): scores = compare_samples(features, db_sample) if scores["overall_similarity"] >= threshold: matches.append({ "sha256": db_sample.get("sha256", ""), "family": db_sample.get("family", "unknown"), "type": db_sample.get("type", "unknown"), "similarity": scores["overall_similarity"], "matching_features": [ k for k, v in scores.items() if k != "overall_similarity" and v > 0.5 ], }) matches.sort(key=lambda x: x["similarity"], reverse=True) matches = matches[:top_k] # Determine classification from top matches if matches: family_votes = Counter(m["family"] for m in matches) best_family = family_votes.most_common(1)[0][0] confidence = matches[0]["similarity"] else: best_family = "unknown" confidence = 0.0 return { "classification": { "family": best_family, "confidence": round(confidence, 4), "method": "feature_similarity", "is_novel": confidence < threshold, }, "similar_samples": matches, } def get_database_stats(database_path: str) -> dict: """Generate statistics from the classification database. Args: database_path: Path to the database JSON. Returns: Dictionary with database statistics. """ db = json.loads(Path(database_path).read_text()) return { "total_samples": db.get("total_samples", 0), "families": db.get("families", {}), "unique_families": len(db.get("families", {})), "created": db.get("created", ""), } # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Malware classification via feature extraction and similarity scoring.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: %(prog)s --extract-features --input sample.exe --output features.json %(prog)s --extract-behavioral --input sandbox_report.json --output behavior.json %(prog)s --compare --input a.exe --reference b.exe --output comparison.json %(prog)s --classify --input sample.exe --database db.json --output result.json %(prog)s --init-db --output classification_db.json %(prog)s --add-to-db --input sample.exe --family emotet --database db.json %(prog)s --db-stats --database classification_db.json """, ) # Mode selection mode = parser.add_mutually_exclusive_group() mode.add_argument("--extract-features", action="store_true", help="Extract static features from PE binary") mode.add_argument("--extract-behavioral", action="store_true", help="Extract behavioral features from sandbox report") mode.add_argument("--compare", action="store_true", help="Compare two samples") mode.add_argument("--cluster", action="store_true", help="Cluster samples from feature file") mode.add_argument("--classify", action="store_true", help="Classify sample against database") mode.add_argument("--init-db", action="store_true", help="Initialize classification database") mode.add_argument("--add-to-db", action="store_true", help="Add classified sample to database") mode.add_argument("--query-db", action="store_true", help="Query database by family") mode.add_argument("--db-stats", action="store_true", help="Show database statistics") # Input/output parser.add_argument("--input", "-i", help="Input file or directory") parser.add_argument("--reference", help="Reference sample for comparison") parser.add_argument("--output", "-o", help="Output JSON file") parser.add_argument("--database", help="Classification database path") # Feature extraction options parser.add_argument("--feature-types", default="imports,sections,header,strings", help="Comma-separated feature types to extract") parser.add_argument("--format", default="json", help="Input format (json, strace)") parser.add_argument("--capa-results", help="Capa JSON results to incorporate as features") parser.add_argument("--yara-results", help="YARA match results to incorporate as features") # Classification options parser.add_argument("--family", help="Malware family name") parser.add_argument("--type", default="unknown", help="Malware type") parser.add_argument("--campaign", default="", help="Campaign identifier") parser.add_argument("--confidence", type=float, default=1.0, help="Classification confidence") parser.add_argument("--threshold", type=float, default=0.5, help="Similarity threshold") parser.add_argument("--top-k", type=int, default=5, help="Number of top matches to return") # Clustering options parser.add_argument("--algorithm", choices=["dbscan", "hierarchical"], default="dbscan", help="Clustering algorithm") parser.add_argument("--eps", type=float, default=0.3, help="DBSCAN eps parameter") parser.add_argument("--min-samples", type=int, default=2, help="DBSCAN min_samples") parser.add_argument("--distance-threshold", type=float, default=0.5, help="Hierarchical clustering distance threshold") parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output") return parser.parse_args() def main() -> None: args = parse_arguments() logging.basicConfig( level=logging.DEBUG if args.verbose else logging.INFO, format="%(levelname)s: %(message)s", ) if args.extract_features: if not args.input: logger.error("--input required") sys.exit(1) input_path = Path(args.input) if input_path.is_dir(): results = [] for f in sorted(input_path.iterdir()): if f.is_file(): try: features = extract_static_features(str(f)) results.append(features) except Exception as e: logger.error(f"Failed on {f.name}: {e}") output = {"samples": results, "total": len(results)} else: output = extract_static_features(args.input) output_json = json.dumps(output, indent=2) if args.output: Path(args.output).write_text(output_json) logger.info(f"Features written to {args.output}") else: print(output_json) elif args.extract_behavioral: if not args.input: logger.error("--input required") sys.exit(1) output = extract_behavioral_features(args.input, args.format) output_json = json.dumps(output, indent=2) if args.output: Path(args.output).write_text(output_json) else: print(output_json) elif args.compare: if not args.input or not args.reference: logger.error("--input and --reference required for --compare") sys.exit(1) features_a = extract_static_features(args.input) features_b = extract_static_features(args.reference) scores = compare_samples(features_a, features_b) output = { "sample_a": features_a.get("hashes", {}).get("sha256", ""), "sample_b": features_b.get("hashes", {}).get("sha256", ""), "similarity_scores": scores, } output_json = json.dumps(output, indent=2) if args.output: Path(args.output).write_text(output_json) else: print(output_json) elif args.classify: if not args.input or not args.database: logger.error("--input and --database required for --classify") sys.exit(1) features = extract_static_features(args.input) result = classify_against_database( features, args.database, top_k=args.top_k, threshold=args.threshold, ) result["sample"] = { "sha256": features["hashes"]["sha256"], "file_name": features["file_name"], "file_size": features["file_size"], } result["features"] = {"static": features} output_json = json.dumps(result, indent=2) if args.output: Path(args.output).write_text(output_json) else: print(output_json) elif args.init_db: if not args.output: logger.error("--output required for --init-db") sys.exit(1) init_database(args.output) elif args.add_to_db: if not all([args.input, args.family, args.database]): logger.error("--input, --family, and --database required") sys.exit(1) features = extract_static_features(args.input) entry = add_to_database( database_path=args.database, features=features, family=args.family, malware_type=args.type, campaign=args.campaign, confidence=args.confidence, ) print(json.dumps(entry, indent=2)) elif args.query_db: if not args.database: logger.error("--database required") sys.exit(1) db = json.loads(Path(args.database).read_text()) results = [s for s in db.get("samples", []) if not args.family or s.get("family") == args.family] print(json.dumps({"results": results, "count": len(results)}, indent=2)) elif args.db_stats: if not args.database: logger.error("--database required") sys.exit(1) stats = get_database_stats(args.database) print(json.dumps(stats, indent=2)) elif args.cluster: logger.info("Clustering requires scikit-learn. Install with: pip install scikit-learn") logger.info("Use: --algorithm dbscan --eps 0.3 --min-samples 2") logger.info("Or: --algorithm hierarchical --distance-threshold 0.5") # Clustering skeleton - requires scikit-learn try: from sklearn.cluster import DBSCAN, AgglomerativeClustering from sklearn.preprocessing import StandardScaler import numpy as np if not args.input: logger.error("--input required (path to features JSON)") sys.exit(1) with open(args.input) as f: data = json.load(f) samples = data.get("samples", [data] if "hashes" in data else []) # Build numeric feature vectors from section entropies and import counts vectors = [] labels = [] for s in samples: vec = s.get("section_entropies", [0.0] * 5)[:5] while len(vec) < 5: vec.append(0.0) vec.append(s.get("import_count", 0)) vec.append(s.get("file_entropy", 0.0)) vec.append(s.get("file_size", 0) / 1e6) vectors.append(vec) labels.append(s.get("hashes", {}).get("sha256", "unknown")) X = StandardScaler().fit_transform(np.array(vectors)) if args.algorithm == "dbscan": model = DBSCAN(eps=args.eps, min_samples=args.min_samples) else: model = AgglomerativeClustering( n_clusters=None, distance_threshold=args.distance_threshold, ) cluster_labels = model.fit_predict(X) clusters = {} for sha, cl in zip(labels, cluster_labels): cl = int(cl) if cl not in clusters: clusters[cl] = [] clusters[cl].append(sha) output = { "algorithm": args.algorithm, "num_clusters": len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0), "noise_samples": int(sum(1 for c in cluster_labels if c == -1)), "clusters": clusters, } output_json = json.dumps(output, indent=2) if args.output: Path(args.output).write_text(output_json) else: print(output_json) except ImportError: logger.error("scikit-learn required for clustering. Install: pip install scikit-learn") sys.exit(1) else: logger.error("Specify a mode. Use --help for options.") sys.exit(1) if __name__ == "__main__": main()