#!/usr/bin/env python3
"""Identify cryptographic algorithms in binary files by scanning for known constants.

Detects AES S-box, RC4 KSA patterns, MD5/SHA initialization vectors, CRC32 tables,
and other well-known cryptographic constants embedded in binary data.
"""

from __future__ import annotations

import argparse
import json
import os
import struct
import sys
from datetime import datetime
from pathlib import Path
from typing import BinaryIO

# AES S-box (first 32 bytes are sufficient for reliable detection)
AES_SBOX_PARTIAL = bytes([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5,
    0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0,
    0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
])

# AES inverse S-box (first 32 bytes)
AES_INV_SBOX_PARTIAL = bytes([
    0x52, 0x09, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38,
    0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB,
    0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87,
    0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB,
])

# AES RCON values
AES_RCON = bytes([0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36])

# MD5 initialization values (as little-endian 32-bit integers)
MD5_INIT = [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476]

# MD5 T-table constants (first 8)
MD5_T_TABLE = [
    0xD76AA478, 0xE8C7B756, 0x242070DB, 0xC1BDCEEE,
    0xF57C0FAF, 0x4787C62A, 0xA8304613, 0xFD469501,
]

# SHA-1 initialization values
SHA1_INIT = [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0]

# SHA-1 round constants
SHA1_K = [0x5A827999, 0x6ED9EBA1, 0x8F1BBCDC, 0xCA62C1D6]

# SHA-256 initialization values
SHA256_INIT = [
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
    0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
]

# SHA-256 round constants (first 8)
SHA256_K = [
    0x428A2F98, 0x71374491, 0xB5C0FBCF, 0xE9B5DBA5,
    0x3956C25B, 0x59F111F1, 0x923F82A4, 0xAB1C5ED5,
]

# CRC32 polynomial and table values
CRC32_POLY = 0xEDB88320
CRC32_TABLE_START = [
    0x00000000, 0x77073096, 0xEE0E612C, 0x990951BA,
    0x076DC419, 0x706AF48F, 0xE963A535, 0x9E6495A3,
]

# Blowfish P-array initial values (first 8)
BLOWFISH_P = [
    0x243F6A88, 0x85A308D3, 0x13198A2E, 0x03707344,
    0xA4093822, 0x299F31D0, 0x082EFA98, 0xEC4E6C89,
]

# RC4 KSA identity permutation pattern (0x00-0xFF sequential)
RC4_IDENTITY = bytes(range(256))

# DES initial permutation table (first 16 values)
DES_IP = bytes([
    58, 50, 42, 34, 26, 18, 10, 2,
    60, 52, 44, 36, 28, 20, 12, 4,
])

# ChaCha20/Salsa20 constant "expand 32-byte k"
CHACHA_CONST = b"expand 32-byte k"
SALSA_CONST = b"expand 32-byte k"


CRYPTO_SIGNATURES = []


def _register_byte_signature(name: str, algo: str, data: bytes, min_match: int = None):
    """Register a byte-pattern signature."""
    CRYPTO_SIGNATURES.append({
        "name": name,
        "algorithm": algo,
        "type": "bytes",
        "pattern": data,
        "min_match": min_match or len(data),
    })


def _register_dword_signature(name: str, algo: str, values: list,
                               endian: str = "both", min_match: int = None):
    """Register a DWORD-pattern signature."""
    CRYPTO_SIGNATURES.append({
        "name": name,
        "algorithm": algo,
        "type": "dwords",
        "values": values,
        "endian": endian,
        "min_match": min_match or len(values),
    })


# Register all signatures
_register_byte_signature("AES S-box", "AES", AES_SBOX_PARTIAL, 16)
_register_byte_signature("AES Inverse S-box", "AES", AES_INV_SBOX_PARTIAL, 16)
_register_byte_signature("AES RCON", "AES", AES_RCON, 8)
_register_byte_signature("DES Initial Permutation", "DES", DES_IP, 16)
_register_byte_signature("ChaCha20/Salsa20 constant", "ChaCha20/Salsa20", CHACHA_CONST, 16)
_register_dword_signature("MD5 Init Vector", "MD5", MD5_INIT, "little", 4)
_register_dword_signature("MD5 T-table", "MD5", MD5_T_TABLE, "little", 4)
_register_dword_signature("SHA-1 Init Vector", "SHA-1", SHA1_INIT, "big", 4)
_register_dword_signature("SHA-1 Round Constants", "SHA-1", SHA1_K, "big", 4)
_register_dword_signature("SHA-256 Init Vector", "SHA-256", SHA256_INIT, "big", 4)
_register_dword_signature("SHA-256 Round Constants", "SHA-256", SHA256_K, "big", 4)
_register_dword_signature("CRC32 Table", "CRC32", CRC32_TABLE_START, "little", 4)
_register_dword_signature("Blowfish P-array", "Blowfish", BLOWFISH_P, "big", 4)


def search_bytes_pattern(data: bytes, pattern: bytes, min_match: int) -> list:
    """Search for a byte pattern in binary data."""
    matches = []
    search_len = min_match
    search_pattern = pattern[:search_len]
    start = 0
    while True:
        offset = data.find(search_pattern, start)
        if offset == -1:
            break
        # Check how much of the full pattern matches
        full_match_len = 0
        for i in range(len(pattern)):
            if offset + i >= len(data):
                break
            if data[offset + i] == pattern[i]:
                full_match_len += 1
            else:
                break
        matches.append({
            "offset": offset,
            "offset_hex": f"0x{offset:08X}",
            "matched_bytes": full_match_len,
            "total_pattern_bytes": len(pattern),
            "confidence": round(full_match_len / len(pattern) * 100, 1),
        })
        start = offset + 1
    return matches


def search_dword_pattern(data: bytes, values: list, endian: str,
                          min_match: int) -> list:
    """Search for DWORD patterns in binary data."""
    matches = []
    endians_to_check = []
    if endian in ("little", "both"):
        endians_to_check.append(("<I", "little-endian"))
    if endian in ("big", "both"):
        endians_to_check.append((">I", "big-endian"))

    for fmt, endian_name in endians_to_check:
        # Build the byte pattern for the first min_match DWORDs
        search_pattern = b""
        for v in values[:min_match]:
            search_pattern += struct.pack(fmt, v)

        start = 0
        while True:
            offset = data.find(search_pattern, start)
            if offset == -1:
                break
            # Count how many consecutive DWORDs match
            matched_count = 0
            for i, v in enumerate(values):
                pos = offset + i * 4
                if pos + 4 > len(data):
                    break
                found_val = struct.unpack(fmt, data[pos:pos + 4])[0]
                if found_val == v:
                    matched_count += 1
                else:
                    break

            matches.append({
                "offset": offset,
                "offset_hex": f"0x{offset:08X}",
                "endian": endian_name,
                "matched_dwords": matched_count,
                "total_pattern_dwords": len(values),
                "confidence": round(matched_count / len(values) * 100, 1),
            })
            start = offset + 1

    return matches


def scan_for_xor_loops(data: bytes) -> list:
    """Detect potential XOR encryption loops in x86 code."""
    xor_patterns = []
    # Common XOR loop patterns in x86:
    # 30 XX - XOR [reg], reg8
    # 31 XX - XOR [reg], reg32
    # 80 3X XX - XOR byte [reg], imm8
    # Pattern: XOR + loop back (E2 = LOOP, 75 = JNZ, 7F = JG)
    for i in range(len(data) - 4):
        if data[i] in (0x30, 0x31, 0x32, 0x33, 0x34, 0x35):
            # Look for a nearby loop instruction within 32 bytes
            for j in range(i + 2, min(i + 32, len(data) - 1)):
                if data[j] in (0xE2, 0x75, 0xEB) and data[j + 1] < 0x80:
                    xor_patterns.append({
                        "offset": i,
                        "offset_hex": f"0x{i:08X}",
                        "xor_opcode": f"0x{data[i]:02X}",
                        "loop_offset": j,
                        "type": "potential_xor_loop",
                    })
                    break
    return xor_patterns[:50]  # Limit results


def scan_for_rc4(data: bytes) -> list:
    """Detect potential RC4 KSA initialization patterns."""
    matches = []
    # Look for the RC4 identity permutation (0x00, 0x01, 0x02, ..., 0xFF)
    # which is created during KSA initialization
    identity_16 = bytes(range(16))
    start = 0
    while True:
        offset = data.find(identity_16, start)
        if offset == -1:
            break
        # Check how far the sequential pattern extends
        seq_len = 0
        for i in range(256):
            if offset + i >= len(data):
                break
            if data[offset + i] == i:
                seq_len += 1
            else:
                break
        if seq_len >= 32:
            matches.append({
                "offset": offset,
                "offset_hex": f"0x{offset:08X}",
                "sequential_bytes": seq_len,
                "likely_rc4_state": seq_len >= 128,
                "confidence": min(round(seq_len / 256 * 100, 1), 100.0),
            })
        start = offset + 1
    return matches


def analyze_binary(binary_path: str, scan_xor: bool = True) -> dict:
    """Perform full cryptographic analysis on a binary file."""
    with open(binary_path, "rb") as f:
        data = f.read()

    results = {
        "binary": os.path.basename(binary_path),
        "binary_path": os.path.abspath(binary_path),
        "file_size": len(data),
        "timestamp": datetime.utcnow().isoformat(),
        "identified_algorithms": [],
        "xor_loops": [],
        "rc4_state_tables": [],
        "summary": {},
    }

    algorithms_found = set()

    for sig in CRYPTO_SIGNATURES:
        if sig["type"] == "bytes":
            matches = search_bytes_pattern(data, sig["pattern"], sig["min_match"])
        else:
            matches = search_dword_pattern(
                data, sig["values"], sig["endian"], sig["min_match"]
            )

        if matches:
            algorithms_found.add(sig["algorithm"])
            results["identified_algorithms"].append({
                "signature": sig["name"],
                "algorithm": sig["algorithm"],
                "matches": matches,
            })

    if scan_xor:
        xor_loops = scan_for_xor_loops(data)
        if xor_loops:
            results["xor_loops"] = xor_loops

    rc4_states = scan_for_rc4(data)
    if rc4_states:
        results["rc4_state_tables"] = rc4_states
        if any(m.get("likely_rc4_state") for m in rc4_states):
            algorithms_found.add("RC4")

    results["summary"] = {
        "algorithms_detected": sorted(algorithms_found),
        "total_signatures_matched": len(results["identified_algorithms"]),
        "xor_loops_found": len(results["xor_loops"]),
        "rc4_states_found": len(results["rc4_state_tables"]),
    }

    return results


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Identify cryptographic algorithms in binary files by scanning for known constants"
    )
    parser.add_argument(
        "--input", "--binary", "-b",
        dest="binary",
        required=True,
        help="Path to the binary file to analyze",
    )
    parser.add_argument(
        "--output", "-o",
        default=None,
        help="Output JSON file path (default: stdout)",
    )
    parser.add_argument(
        "--no-xor-scan",
        action="store_true",
        help="Skip XOR loop detection (faster analysis)",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Print detailed progress information",
    )
    parser.add_argument(
        "--format", default="json", choices=["json", "text", "csv"],
        help="Output format (default: json)",
    )

    args = parser.parse_args()

    if not os.path.isfile(args.binary):
        print(f"[!] Error: File not found: {args.binary}", file=sys.stderr)
        sys.exit(1)

    if args.verbose:
        print(f"[*] Analyzing: {args.binary}")
        print(f"[*] File size: {os.path.getsize(args.binary)} bytes")

    try:
        results = analyze_binary(args.binary, scan_xor=not args.no_xor_scan)
    except PermissionError:
        print(f"[!] Permission denied: {args.binary}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"[!] Error analyzing binary: {e}", file=sys.stderr)
        sys.exit(1)

    output_json = json.dumps(results, indent=2)

    if args.output:
        with open(args.output, "w") as f:
            f.write(output_json)
        print(f"[+] Results saved to: {args.output}")
    else:
        print(output_json)

    # Print summary
    summary = results["summary"]
    if summary["algorithms_detected"]:
        print(f"\n[+] Algorithms detected: {', '.join(summary['algorithms_detected'])}")
    else:
        print("\n[-] No known cryptographic constants detected")

    if summary["xor_loops_found"]:
        print(f"[+] Potential XOR loops: {summary['xor_loops_found']}")
    if summary["rc4_states_found"]:
        print(f"[+] RC4 state tables: {summary['rc4_states_found']}")


if __name__ == "__main__":
    main()