#!/usr/bin/env python3
"""Analyze cryptographic implementations in ransomware samples.

Examines a ransomware binary for crypto-related constants, API imports,
and patterns.  Optionally compares an encrypted file against its original
to determine encryption characteristics (block size, mode, partial vs full).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from pathlib import Path
from typing import Any


# Well-known cryptographic constants for detection
_CRYPTO_SIGNATURES: list[dict[str, Any]] = [
    {
        "name": "AES S-box",
        "pattern": b"\x63\x7c\x77\x7b\xf2\x6b\x6f\xc5",
        "algorithm": "AES",
    },
    {
        "name": "AES inverse S-box",
        "pattern": b"\x52\x09\x6a\xd5\x30\x36\xa5\x38",
        "algorithm": "AES",
    },
    {
        "name": "ChaCha20/Salsa20 constant",
        "pattern": b"expand 32-byte k",
        "algorithm": "ChaCha20/Salsa20",
    },
    {
        "name": "ChaCha20/Salsa20 constant (16-byte key)",
        "pattern": b"expand 16-byte k",
        "algorithm": "ChaCha20/Salsa20",
    },
    {
        "name": "SHA-256 initial hash value",
        "pattern": b"\x6a\x09\xe6\x67\xbb\x67\xae\x85",
        "algorithm": "SHA-256",
    },
    {
        "name": "RSA public exponent (65537 LE)",
        "pattern": b"\x01\x00\x01\x00",
        "algorithm": "RSA",
    },
]

# Windows CryptoAPI and BCrypt function names
_CRYPTO_API_NAMES: list[str] = [
    "CryptAcquireContext", "CryptGenRandom", "CryptEncrypt", "CryptDecrypt",
    "CryptDeriveKey", "CryptImportKey", "CryptExportKey", "CryptGenKey",
    "CryptCreateHash", "CryptHashData", "CryptDestroyKey",
    "BCryptOpenAlgorithmProvider", "BCryptGenerateSymmetricKey",
    "BCryptEncrypt", "BCryptDecrypt", "BCryptGenRandom",
    "BCryptDeriveKeyPBKDF2", "BCryptGenerateKeyPair",
]


def _sha256(path: Path) -> str:
    """Compute SHA-256 hash of a file."""
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def _scan_crypto_constants(data: bytes) -> list[dict[str, Any]]:
    """Scan binary data for known cryptographic constants."""
    findings: list[dict[str, Any]] = []
    for sig in _CRYPTO_SIGNATURES:
        offset = data.find(sig["pattern"])
        if offset != -1:
            findings.append({
                "constant": sig["name"],
                "algorithm": sig["algorithm"],
                "offset": hex(offset),
            })
    return findings


def _scan_crypto_apis(data: bytes) -> list[str]:
    """Scan for references to cryptographic API function names."""
    found: list[str] = []
    for api in _CRYPTO_API_NAMES:
        if api.encode("ascii") in data or api.encode("utf-16-le") in data:
            found.append(api)
    return found


def _extract_strings(data: bytes, min_length: int = 6) -> list[str]:
    """Extract printable ASCII strings from binary data."""
    pattern = re.compile(rb"[\x20-\x7e]{%d,}" % min_length)
    return [m.decode("ascii", errors="replace") for m in pattern.findall(data)]


def _analyze_encryption_pattern(
    original: Path, encrypted: Path
) -> dict[str, Any]:
    """Compare an original file with its encrypted version to infer crypto details."""
    orig_data = original.read_bytes()
    enc_data = encrypted.read_bytes()

    result: dict[str, Any] = {
        "original_size": len(orig_data),
        "encrypted_size": len(enc_data),
        "size_difference": len(enc_data) - len(orig_data),
        "full_file_encrypted": True,
        "header_preserved": False,
        "footer_appended": False,
    }

    # Check if the file sizes differ (appended key material, IV, etc.)
    if len(enc_data) > len(orig_data):
        result["footer_appended"] = True
        result["appended_bytes"] = len(enc_data) - len(orig_data)

    # Check if the first bytes are plaintext (partial/intermittent encryption)
    if len(orig_data) >= 16 and len(enc_data) >= 16:
        if orig_data[:8] == enc_data[:8]:
            result["full_file_encrypted"] = False
            result["header_preserved"] = True

    # Estimate block size by looking for repeating patterns in ECB mode
    block_sizes_to_check = [16, 32, 8]
    for bs in block_sizes_to_check:
        blocks = [enc_data[i:i + bs] for i in range(0, min(len(enc_data), 4096), bs)]
        unique = len(set(blocks))
        if unique < len(blocks) * 0.9 and len(blocks) > 2:
            result["possible_ecb_mode"] = True
            result["estimated_block_size"] = bs
            break

    return result


def analyze_sample(
    sample_path: Path,
    encrypted_file: Path | None = None,
    original_file: Path | None = None,
) -> dict[str, Any]:
    """Perform cryptographic analysis of a ransomware sample."""
    data = sample_path.read_bytes()

    result: dict[str, Any] = {
        "sample": {
            "filename": sample_path.name,
            "sha256": _sha256(sample_path),
            "file_size": len(data),
        },
        "crypto_constants": _scan_crypto_constants(data),
        "crypto_apis": _scan_crypto_apis(data),
        "algorithms_detected": [],
        "encryption_pattern": None,
        "assessment": {
            "key_generation": "unknown",
            "weaknesses_found": [],
            "recovery_feasible": None,
        },
    }

    # Summarise detected algorithms
    algos = {c["algorithm"] for c in result["crypto_constants"]}
    result["algorithms_detected"] = sorted(algos)

    # Check for weak patterns in strings
    strings = _extract_strings(data)
    weak_indicators = []
    for s in strings:
        if re.search(r"srand|rand\(\)|time\(0\)|GetTickCount", s, re.IGNORECASE):
            weak_indicators.append(f"Possible weak RNG: {s[:80]}")
        if re.search(r"ECB", s, re.IGNORECASE):
            weak_indicators.append(f"ECB mode reference: {s[:80]}")

    result["assessment"]["weaknesses_found"] = weak_indicators
    if weak_indicators:
        result["assessment"]["recovery_feasible"] = "possible"

    # Compare encrypted vs original if provided
    if encrypted_file and original_file:
        if encrypted_file.exists() and original_file.exists():
            result["encryption_pattern"] = _analyze_encryption_pattern(
                original_file, encrypted_file
            )

    return result


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Analyze a ransomware sample and return results."""
    return analyze_sample(input_path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Analyze cryptographic implementations in ransomware samples."
    )
    parser.add_argument("--input", type=Path,
                        help="Path to ransomware sample (alias for --sample)")
    parser.add_argument("--sample", type=Path,
                        help="Path to ransomware sample binary")
    parser.add_argument("--encrypted-file", type=Path,
                        help="Path to an encrypted file for pattern analysis")
    parser.add_argument("--original-file", type=Path,
                        help="Path to the original (pre-encryption) version of the file")
    parser.add_argument("--output", type=Path,
                        help="Path to output file (stdout if omitted)")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")
    args = parser.parse_args()

    sample = args.sample or args.input
    if not sample:
        parser.error("Provide a sample via --sample or --input")

    if not sample.exists():
        print(f"[error] file not found: {sample}", file=sys.stderr)
        sys.exit(1)

    result = analyze_sample(
        sample,
        encrypted_file=args.encrypted_file,
        original_file=args.original_file,
    )
    rendered = json.dumps(result, indent=2)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
        print(f"[+] Crypto analysis written to {args.output}", file=sys.stderr)
    else:
        print(rendered)


if __name__ == "__main__":
    main()