#!/usr/bin/env python3
"""
Malware code similarity analyzer.

Compares malware samples using multiple similarity metrics: ssdeep fuzzy hashing,
TLSH locality-sensitive hashing, import table Jaccard similarity, and string
overlap analysis. Outputs pairwise similarity scores and a similarity matrix.

Usage:
    python3 similarity_analyzer.py --samples s1.exe s2.exe --metrics all
    python3 similarity_analyzer.py --samples-dir ./samples/ --output matrix.json
    python3 similarity_analyzer.py --samples s1.exe s2.exe --metrics ssdeep imports
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import re
import string
import struct
import sys
from datetime import datetime, timezone
from pathlib import Path

# ---------------------------------------------------------------------------
# Optional dependency imports with graceful fallback
# ---------------------------------------------------------------------------

try:
    import ssdeep as _ssdeep
    HAS_SSDEEP = True
except ImportError:
    HAS_SSDEEP = False

try:
    import tlsh as _tlsh
    HAS_TLSH = True
except ImportError:
    HAS_TLSH = False

try:
    import pefile
    HAS_PEFILE = True
except ImportError:
    HAS_PEFILE = False

# ---------------------------------------------------------------------------
# String extraction (standalone, no external dependency)
# ---------------------------------------------------------------------------

PRINTABLE = set(string.printable) - set(string.whitespace) | {' '}


def extract_strings(file_path: str, min_length: int = 6, max_strings: int = 5000) -> list:
    """Extract printable ASCII strings from a binary file."""
    strings_found = []
    try:
        with open(file_path, "rb") as f:
            data = f.read()
    except OSError:
        return []

    current = []
    for byte in data:
        ch = chr(byte)
        if ch in PRINTABLE:
            current.append(ch)
        else:
            if len(current) >= min_length:
                strings_found.append("".join(current))
                if len(strings_found) >= max_strings:
                    break
            current = []

    if len(current) >= min_length and len(strings_found) < max_strings:
        strings_found.append("".join(current))

    return strings_found


# ---------------------------------------------------------------------------
# Import table extraction
# ---------------------------------------------------------------------------


def extract_imports(file_path: str) -> set:
    """Extract import function names from a PE file."""
    if not HAS_PEFILE:
        return set()

    try:
        pe = pefile.PE(file_path, fast_load=True)
        pe.parse_data_directories(
            directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"]]
        )
    except Exception:
        return set()

    imports = set()
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode("utf-8", errors="replace").lower()
            for imp in entry.imports:
                if imp.name:
                    func_name = imp.name.decode("utf-8", errors="replace")
                    imports.add(f"{dll_name}:{func_name}")
    pe.close()
    return imports


# ---------------------------------------------------------------------------
# Similarity metrics
# ---------------------------------------------------------------------------


def jaccard_similarity(set_a: set, set_b: set) -> float:
    """Compute Jaccard similarity coefficient between two sets."""
    if not set_a and not set_b:
        return 0.0
    intersection = set_a & set_b
    union = set_a | set_b
    return len(intersection) / len(union) if union else 0.0


def compute_ssdeep_hash(file_path: str) -> str:
    """Compute ssdeep fuzzy hash for a file."""
    if not HAS_SSDEEP:
        return None
    try:
        return _ssdeep.hash_from_file(file_path)
    except Exception:
        return None


def compare_ssdeep(hash1: str, hash2: str) -> int:
    """Compare two ssdeep hashes. Returns score 0-100."""
    if not HAS_SSDEEP or not hash1 or not hash2:
        return -1
    try:
        return _ssdeep.compare(hash1, hash2)
    except Exception:
        return -1


def compute_tlsh_hash(file_path: str) -> str:
    """Compute TLSH hash for a file."""
    if not HAS_TLSH:
        return None
    try:
        with open(file_path, "rb") as f:
            data = f.read()
        h = _tlsh.hash(data)
        return h if h else None
    except Exception:
        return None


def compare_tlsh(hash1: str, hash2: str) -> int:
    """Compare two TLSH hashes. Returns distance (lower = more similar)."""
    if not HAS_TLSH or not hash1 or not hash2:
        return -1
    try:
        return _tlsh.diff(hash1, hash2)
    except Exception:
        return -1


def compare_imports(file1: str, file2: str) -> dict:
    """Compare import tables of two PE files using Jaccard similarity."""
    imports1 = extract_imports(file1)
    imports2 = extract_imports(file2)

    if not imports1 and not imports2:
        return {
            "jaccard": 0.0,
            "shared_count": 0,
            "total_union": 0,
            "shared_imports": [],
            "available": False,
            "reason": "No imports extracted (not PE files or pefile not installed)",
        }

    shared = imports1 & imports2
    union = imports1 | imports2

    return {
        "jaccard": len(shared) / len(union) if union else 0.0,
        "shared_count": len(shared),
        "total_file1": len(imports1),
        "total_file2": len(imports2),
        "total_union": len(union),
        "shared_imports": sorted(list(shared))[:50],  # Limit output size
        "available": True,
    }


def compare_strings(file1: str, file2: str, min_length: int = 8) -> dict:
    """Compare extracted strings between two files using Jaccard similarity."""
    strings1 = set(extract_strings(file1, min_length=min_length))
    strings2 = set(extract_strings(file2, min_length=min_length))

    if not strings1 and not strings2:
        return {
            "jaccard": 0.0,
            "shared_count": 0,
            "total_union": 0,
            "shared_strings": [],
        }

    shared = strings1 & strings2
    union = strings1 | strings2

    # Filter out very common strings (reduce noise)
    common_noise = {
        "This program cannot be run in DOS mode",
        "Rich",
        ".text",
        ".data",
        ".rdata",
        ".rsrc",
        ".reloc",
    }
    meaningful_shared = shared - common_noise

    return {
        "jaccard": len(shared) / len(union) if union else 0.0,
        "shared_count": len(shared),
        "meaningful_shared_count": len(meaningful_shared),
        "total_file1": len(strings1),
        "total_file2": len(strings2),
        "total_union": len(union),
        "shared_strings": sorted(list(meaningful_shared))[:100],  # Limit output
    }


# ---------------------------------------------------------------------------
# Composite similarity
# ---------------------------------------------------------------------------


def compute_sha256(file_path: str) -> str:
    """Compute SHA256 hash of a file."""
    sha256 = hashlib.sha256()
    try:
        with open(file_path, "rb") as f:
            for block in iter(lambda: f.read(65536), b""):
                sha256.update(block)
        return sha256.hexdigest()
    except OSError:
        return None


def compare_samples(file1: str, file2: str, metrics: list) -> dict:
    """Compare two samples across specified metrics."""
    result = {
        "file1": os.path.basename(file1),
        "file2": os.path.basename(file2),
        "metrics": {},
    }

    scores = []

    if "ssdeep" in metrics:
        h1 = compute_ssdeep_hash(file1)
        h2 = compute_ssdeep_hash(file2)
        score = compare_ssdeep(h1, h2)
        result["metrics"]["ssdeep"] = {
            "hash1": h1,
            "hash2": h2,
            "score": score,
            "normalized": score / 100.0 if score >= 0 else None,
            "available": HAS_SSDEEP and score >= 0,
        }
        if score >= 0:
            scores.append(score / 100.0)

    if "tlsh" in metrics:
        h1 = compute_tlsh_hash(file1)
        h2 = compute_tlsh_hash(file2)
        distance = compare_tlsh(h1, h2)
        # Normalize TLSH distance to 0-1 similarity (rough approximation)
        normalized = max(0.0, 1.0 - (distance / 300.0)) if distance >= 0 else None
        result["metrics"]["tlsh"] = {
            "hash1": h1,
            "hash2": h2,
            "distance": distance,
            "normalized": normalized,
            "available": HAS_TLSH and distance >= 0,
        }
        if normalized is not None:
            scores.append(normalized)

    if "imports" in metrics:
        import_result = compare_imports(file1, file2)
        result["metrics"]["imports"] = import_result
        if import_result.get("available", False):
            scores.append(import_result["jaccard"])

    if "strings" in metrics:
        string_result = compare_strings(file1, file2)
        result["metrics"]["strings"] = string_result
        if string_result["total_union"] > 0:
            scores.append(string_result["jaccard"])

    # Compute overall similarity (average of available metrics)
    if scores:
        result["overall_similarity"] = sum(scores) / len(scores)
    else:
        result["overall_similarity"] = None

    # Classification
    overall = result["overall_similarity"]
    if overall is not None:
        if overall >= 0.9:
            result["classification"] = "Same sample or recompilation"
        elif overall >= 0.7:
            result["classification"] = "Same family variant - high confidence"
        elif overall >= 0.5:
            result["classification"] = "Related family - medium confidence"
        elif overall >= 0.3:
            result["classification"] = "Possible connection - low confidence"
        else:
            result["classification"] = "Likely unrelated"
    else:
        result["classification"] = "Unable to determine (no metrics available)"

    return result


# ---------------------------------------------------------------------------
# Matrix generation
# ---------------------------------------------------------------------------


def build_similarity_matrix(sample_paths: list, metrics: list) -> dict:
    """Build a pairwise similarity matrix for all samples."""
    n = len(sample_paths)
    names = [os.path.basename(p) for p in sample_paths]

    # Precompute hashes for efficiency
    hashes = {}
    for path in sample_paths:
        name = os.path.basename(path)
        hashes[name] = {
            "sha256": compute_sha256(path),
            "path": path,
        }

    # Build pairwise comparisons
    comparisons = []
    matrix = [[0.0] * n for _ in range(n)]

    for i in range(n):
        matrix[i][i] = 1.0  # Self-similarity
        for j in range(i + 1, n):
            result = compare_samples(sample_paths[i], sample_paths[j], metrics)
            comparisons.append(result)

            overall = result.get("overall_similarity")
            if overall is not None:
                matrix[i][j] = overall
                matrix[j][i] = overall

    return {
        "timestamp": datetime.now(tz=timezone.utc).isoformat(),
        "sample_count": n,
        "sample_names": names,
        "sample_hashes": hashes,
        "metrics_used": metrics,
        "matrix": matrix,
        "pairwise_comparisons": comparisons,
    }


def format_matrix_csv(matrix_data: dict) -> str:
    """Format similarity matrix as CSV."""
    names = matrix_data["sample_names"]
    matrix = matrix_data["matrix"]

    lines = ["," + ",".join(names)]
    for i, name in enumerate(names):
        row_values = [f"{matrix[i][j]:.3f}" for j in range(len(names))]
        lines.append(f"{name},{','.join(row_values)}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def collect_samples(args) -> list:
    """Collect sample file paths from CLI arguments."""
    samples = []

    if args.samples:
        for path in args.samples:
            if os.path.isfile(path):
                samples.append(os.path.abspath(path))
            else:
                print(f"Warning: File not found: {path}", file=sys.stderr)

    if args.samples_dir:
        if os.path.isdir(args.samples_dir):
            for entry in sorted(os.listdir(args.samples_dir)):
                full_path = os.path.join(args.samples_dir, entry)
                if os.path.isfile(full_path):
                    samples.append(os.path.abspath(full_path))
        else:
            print(f"Error: Directory not found: {args.samples_dir}", file=sys.stderr)
            sys.exit(1)

    return samples


def resolve_metrics(metrics_arg: list) -> list:
    """Resolve metric names, expanding 'all'."""
    all_metrics = ["ssdeep", "tlsh", "imports", "strings"]
    if "all" in metrics_arg:
        return all_metrics
    valid = []
    for m in metrics_arg:
        if m in all_metrics:
            valid.append(m)
        else:
            print(f"Warning: Unknown metric '{m}', skipping", file=sys.stderr)
    return valid if valid else all_metrics


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Analyze code similarity between malware samples",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  %(prog)s --samples s1.exe s2.exe --metrics all\n"
            "  %(prog)s --samples-dir ./samples/ --output matrix.json\n"
            "  %(prog)s --samples s1.exe s2.exe --metrics ssdeep imports --verbose\n"
        ),
    )
    parser.add_argument(
        "--input", "--samples", "-s", nargs="+",
        dest="samples",
        help="Paths to sample files to compare",
    )
    parser.add_argument(
        "--samples-dir", "-d",
        help="Directory containing samples to compare",
    )
    parser.add_argument(
        "--metrics", "-m", nargs="+", default=["all"],
        help="Similarity metrics to use: ssdeep, tlsh, imports, strings, all (default: all)",
    )
    parser.add_argument(
        "--output", "-o",
        help="Output file path (default: stdout)",
    )
    parser.add_argument(
        "--format", "-f", choices=["json", "csv", "matrix"],
        default="json",
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true",
        help="Include detailed shared artifacts in output",
    )
    parser.add_argument(
        "--min-string-length", type=int, default=8,
        help="Minimum string length for string comparison (default: 8)",
    )

    args = parser.parse_args()

    # Collect samples
    samples = collect_samples(args)
    if len(samples) < 2:
        print("Error: At least 2 samples are required for comparison", file=sys.stderr)
        sys.exit(1)

    metrics = resolve_metrics(args.metrics)

    # Report available libraries
    print(f"Samples: {len(samples)}", file=sys.stderr)
    print(f"Metrics: {', '.join(metrics)}", file=sys.stderr)
    if "ssdeep" in metrics and not HAS_SSDEEP:
        print("Warning: ssdeep not installed, skipping ssdeep metric", file=sys.stderr)
    if "tlsh" in metrics and not HAS_TLSH:
        print("Warning: tlsh not installed, skipping tlsh metric", file=sys.stderr)
    if "imports" in metrics and not HAS_PEFILE:
        print("Warning: pefile not installed, skipping imports metric", file=sys.stderr)

    # Compare
    if len(samples) == 2:
        result = compare_samples(samples[0], samples[1], metrics)
        output_data = result
    else:
        output_data = build_similarity_matrix(samples, metrics)

    # Format output
    if args.format == "csv" and isinstance(output_data, dict) and "matrix" in output_data:
        output_text = format_matrix_csv(output_data)
    else:
        output_text = json.dumps(output_data, indent=2, default=str)

    # Write output
    if args.output:
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        Path(args.output).write_text(output_text, encoding="utf-8")
        print(f"Results written to {args.output}", file=sys.stderr)
    else:
        print(output_text)


if __name__ == "__main__":
    main()