#!/usr/bin/env python3 """Manage malware sample inventory with SQLite database. Track samples, metadata, analysis status, and generate reports. """ from __future__ import annotations import argparse import hashlib import json import os import sqlite3 import sys from datetime import datetime from pathlib import Path DB_SCHEMA = """ CREATE TABLE IF NOT EXISTS samples ( id INTEGER PRIMARY KEY AUTOINCREMENT, sha256 TEXT UNIQUE NOT NULL, md5 TEXT, file_name TEXT, file_size INTEGER, file_type TEXT, family TEXT, source TEXT, notes TEXT, status TEXT DEFAULT 'new', storage_path TEXT, acquired_at TEXT, analyzed_at TEXT, created_at TEXT DEFAULT CURRENT_TIMESTAMP ); """ def init_db(db_path) -> dict: """Initialize the sample database.""" conn = sqlite3.connect(db_path) conn.execute(DB_SCHEMA) conn.commit() return conn def compute_hashes(filepath) -> tuple: """Compute MD5 and SHA256 hashes.""" md5 = hashlib.md5() sha256 = hashlib.sha256() with open(filepath, "rb") as f: for chunk in iter(lambda: f.read(8192), b""): md5.update(chunk) sha256.update(chunk) return md5.hexdigest(), sha256.hexdigest() def add_sample(conn, filepath, family=None, source=None, notes=None) -> bool: """Add a sample to the database.""" path = Path(filepath) if not path.exists(): print(f"[!] File not found: {filepath}", file=sys.stderr) return False md5, sha256 = compute_hashes(filepath) file_size = path.stat().st_size try: conn.execute( """INSERT INTO samples (sha256, md5, file_name, file_size, family, source, notes, status, storage_path, acquired_at) VALUES (?, ?, ?, ?, ?, ?, ?, 'new', ?, ?)""", (sha256, md5, path.name, file_size, family, source, notes, str(path.absolute()), datetime.now().isoformat()) ) conn.commit() print(f"[+] Added: {sha256[:16]}... ({path.name})") return True except sqlite3.IntegrityError: print(f"[!] Sample already exists: {sha256[:16]}...") return False def search_samples(conn, family=None, status=None, sha256=None) -> dict: """Search the sample database.""" query = "SELECT * FROM samples WHERE 1=1" params = [] if family: query += " AND family LIKE ?" params.append(f"%{family}%") if status: query += " AND status = ?" params.append(status) if sha256: query += " AND sha256 LIKE ?" params.append(f"{sha256}%") cursor = conn.execute(query, params) columns = [desc[0] for desc in cursor.description] return [dict(zip(columns, row)) for row in cursor.fetchall()] def update_status(conn, sha256, status) -> None: """Update sample analysis status.""" conn.execute( "UPDATE samples SET status = ?, analyzed_at = ? WHERE sha256 LIKE ?", (status, datetime.now().isoformat() if status == "analyzed" else None, f"{sha256}%") ) conn.commit() print(f"[+] Updated {sha256[:16]}... → {status}") def generate_inventory(conn) -> dict: """Generate inventory report.""" cursor = conn.execute(""" SELECT COUNT(*) as total, COUNT(CASE WHEN status = 'new' THEN 1 END) as new_count, COUNT(CASE WHEN status = 'analyzing' THEN 1 END) as analyzing, COUNT(CASE WHEN status = 'analyzed' THEN 1 END) as analyzed, SUM(file_size) as total_size FROM samples """) stats = cursor.fetchone() cursor = conn.execute(""" SELECT family, COUNT(*) as count FROM samples WHERE family IS NOT NULL GROUP BY family ORDER BY count DESC LIMIT 20 """) families = cursor.fetchall() cursor = conn.execute(""" SELECT source, COUNT(*) as count FROM samples WHERE source IS NOT NULL GROUP BY source ORDER BY count DESC """) sources = cursor.fetchall() report = { "timestamp": datetime.now().isoformat(), "total_samples": stats[0], "status": { "new": stats[1], "analyzing": stats[2], "analyzed": stats[3], }, "total_size_bytes": stats[4] or 0, "top_families": [{"family": f[0], "count": f[1]} for f in families], "sources": [{"source": s[0], "count": s[1]} for s in sources], } return report def main() -> None: parser = argparse.ArgumentParser(description="Manage malware sample inventory") parser.add_argument("--db", default="./samples/database.sqlite", help="Database path") group = parser.add_mutually_exclusive_group(required=True) group.add_argument("--input", "--add", help="Add a sample file") group.add_argument("--search", action="store_true", help="Search samples") group.add_argument("--update-status", nargs=2, metavar=("SHA256", "STATUS"), help="Update sample status") group.add_argument("--inventory", action="store_true", help="Generate inventory report") parser.add_argument("--family", help="Malware family name") parser.add_argument("--source", help="Sample source") parser.add_argument("--notes", help="Notes about the sample") parser.add_argument("--status", help="Filter by status (for search)") parser.add_argument("--sha256", help="Filter by SHA256 prefix (for search)") parser.add_argument("--format", choices=["json", "text"], default="text") args = parser.parse_args() # Ensure database directory exists db_path = Path(args.db) db_path.parent.mkdir(parents=True, exist_ok=True) conn = init_db(args.db) if args.add: add_sample(conn, args.add, args.family, args.source, args.notes) elif args.search: results = search_samples(conn, args.family, args.status, args.sha256) if args.format == "json": print(json.dumps(results, indent=2)) else: print(f"Found {len(results)} samples:\n") for s in results: print(f" {s['sha256'][:16]}... | {s['file_name']} | {s['family'] or 'unknown'} | {s['status']}") elif args.update_status: sha256, status = args.update_status valid_statuses = ["new", "analyzing", "analyzed", "archived"] if status not in valid_statuses: print(f"[!] Invalid status. Use: {', '.join(valid_statuses)}") sys.exit(1) update_status(conn, sha256, status) elif args.inventory: report = generate_inventory(conn) if args.format == "json": print(json.dumps(report, indent=2)) else: print(f"=== Sample Inventory ===") print(f"Total: {report['total_samples']} samples") print(f"Size: {report['total_size_bytes'] / 1024 / 1024:.1f} MB") print(f"Status: New={report['status']['new']}, " f"Analyzing={report['status']['analyzing']}, " f"Analyzed={report['status']['analyzed']}") if report["top_families"]: print(f"\nTop Families:") for f in report["top_families"]: print(f" {f['family']}: {f['count']}") conn.close() if __name__ == "__main__": main()