#!/usr/bin/env python3
"""Framework for decrypting obfuscated strings found in malware samples.

Implements common decryption patterns: XOR (single-byte, multi-byte, rolling),
Base64 (standard and custom alphabet), RC4, AES-CBC, and brute-force XOR.

Usage:
    python3 string_decryptor.py --method xor-single --key 0x5A --input encrypted.bin
    python3 string_decryptor.py --method xor-multi --key "secret" --input encrypted.bin
    python3 string_decryptor.py --method xor-brute --input encrypted.bin
    python3 string_decryptor.py --method base64 --input encoded.txt
    python3 string_decryptor.py --method base64-custom --alphabet "..." --input data.bin
    python3 string_decryptor.py --method rc4 --key "key" --input encrypted.bin
    python3 string_decryptor.py --method aes-cbc --key <hex> --iv <hex> --input encrypted.bin
    python3 string_decryptor.py --method xor-rolling --key 0x41 --input encrypted.bin
"""

from __future__ import annotations

import argparse
import base64
import os
import string
import struct
import sys
from pathlib import Path
from typing import List, Optional, Tuple


def xor_single_byte(data: bytes, key: int) -> bytes:
    """Decrypt data using single-byte XOR."""
    return bytes(b ^ key for b in data)


def xor_multi_byte(data: bytes, key: bytes) -> bytes:
    """Decrypt data using multi-byte (repeating key) XOR."""
    key_len = len(key)
    return bytes(data[i] ^ key[i % key_len] for i in range(len(data)))


def xor_rolling(data: bytes, initial_key: int) -> bytes:
    """Decrypt data using rolling XOR (each byte's key depends on the previous).

    Common pattern: key[i+1] = (key[i] + plaintext[i]) & 0xFF
    """
    result = bytearray()
    key = initial_key & 0xFF
    for b in data:
        decrypted = b ^ key
        result.append(decrypted)
        key = (key + decrypted) & 0xFF
    return bytes(result)


def xor_rolling_variant2(data: bytes, initial_key: int) -> bytes:
    """Rolling XOR variant: key[i+1] = (key[i] + ciphertext[i]) & 0xFF."""
    result = bytearray()
    key = initial_key & 0xFF
    for b in data:
        decrypted = b ^ key
        result.append(decrypted)
        key = (key + b) & 0xFF
    return bytes(result)


def xor_rolling_variant3(data: bytes, initial_key: int) -> bytes:
    """Rolling XOR variant: key rotates (ROL by 1 each iteration)."""
    result = bytearray()
    key = initial_key & 0xFF
    for b in data:
        decrypted = b ^ key
        result.append(decrypted)
        key = ((key << 1) | (key >> 7)) & 0xFF
    return bytes(result)


def base64_decode(data: bytes) -> bytes:
    """Decode standard Base64."""
    # Handle data that might have whitespace or newlines
    cleaned = data.replace(b"\r", b"").replace(b"\n", b"").strip()
    # Add padding if needed
    padding = 4 - (len(cleaned) % 4) if len(cleaned) % 4 else 0
    cleaned += b"=" * padding
    return base64.b64decode(cleaned)


def base64_custom_decode(data: bytes, custom_alphabet: str) -> bytes:
    """Decode Base64 with a custom alphabet."""
    standard = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    if len(custom_alphabet) != 64:
        raise ValueError(f"Custom alphabet must be exactly 64 characters, got {len(custom_alphabet)}")

    # Build translation table
    cleaned = data.replace(b"\r", b"").replace(b"\n", b"").strip()
    translated = cleaned.decode("ascii", errors="replace")
    result = ""
    for ch in translated:
        if ch == "=":
            result += "="
        elif ch in custom_alphabet:
            idx = custom_alphabet.index(ch)
            result += standard[idx]
        else:
            result += ch

    padding = 4 - (len(result) % 4) if len(result) % 4 else 0
    result += "=" * padding
    return base64.b64decode(result)


def rc4_decrypt(data: bytes, key: bytes) -> bytes:
    """Decrypt data using RC4 (ARC4) algorithm."""
    # Key Scheduling Algorithm (KSA)
    S = list(range(256))
    j = 0
    for i in range(256):
        j = (j + S[i] + key[i % len(key)]) & 0xFF
        S[i], S[j] = S[j], S[i]

    # Pseudo-Random Generation Algorithm (PRGA)
    result = bytearray()
    i = 0
    j = 0
    for byte in data:
        i = (i + 1) & 0xFF
        j = (j + S[i]) & 0xFF
        S[i], S[j] = S[j], S[i]
        k = S[(S[i] + S[j]) & 0xFF]
        result.append(byte ^ k)
    return bytes(result)


def aes_cbc_decrypt(data: bytes, key: bytes, iv: bytes) -> bytes:
    """Decrypt data using AES-CBC.

    Tries to use the cryptography library, falls back to PyCryptodome,
    then to a pure-Python fallback warning.
    """
    try:
        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
        from cryptography.hazmat.backends import default_backend
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
        decryptor = cipher.decryptor()
        padded = decryptor.update(data) + decryptor.finalize()
        # Remove PKCS7 padding
        pad_len = padded[-1]
        if pad_len <= 16 and all(b == pad_len for b in padded[-pad_len:]):
            return padded[:-pad_len]
        return padded
    except ImportError:
        pass

    try:
        from Crypto.Cipher import AES
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded = cipher.decrypt(data)
        pad_len = padded[-1]
        if pad_len <= 16 and all(b == pad_len for b in padded[-pad_len:]):
            return padded[:-pad_len]
        return padded
    except ImportError:
        pass

    print("Warning: No AES library available. Install 'cryptography' or 'pycryptodome'.",
          file=sys.stderr)
    print("  pip install cryptography", file=sys.stderr)
    sys.exit(1)


def is_printable_ratio(data: bytes, threshold: float = 0.7) -> bool:
    """Check if a sufficient ratio of bytes are printable ASCII."""
    if not data:
        return False
    printable_count = sum(1 for b in data if 0x20 <= b <= 0x7E or b in (0x0A, 0x0D, 0x09))
    return printable_count / len(data) >= threshold


def extract_printable_strings(data: bytes, min_length: int = 4) -> List[str]:
    """Extract printable ASCII strings from binary data."""
    strings_found = []
    current = []
    for b in data:
        if 0x20 <= b <= 0x7E:
            current.append(chr(b))
        else:
            if len(current) >= min_length:
                strings_found.append("".join(current))
            current = []
    if len(current) >= min_length:
        strings_found.append("".join(current))
    return strings_found


def xor_brute_force(data: bytes, min_string_len: int = 6) -> List[Tuple[int, float, List[str]]]:
    """Brute-force all 256 single-byte XOR keys and rank by printable content.

    Returns list of (key, printable_ratio, found_strings) sorted by score.
    """
    results = []
    for key in range(256):
        if key == 0:
            continue  # Skip identity XOR
        decrypted = xor_single_byte(data, key)
        printable_count = sum(1 for b in decrypted if 0x20 <= b <= 0x7E)
        ratio = printable_count / len(decrypted) if decrypted else 0
        strings = extract_printable_strings(decrypted, min_string_len)
        if strings:
            results.append((key, ratio, strings))

    results.sort(key=lambda x: x[1], reverse=True)
    return results[:10]  # Top 10 candidates


def parse_key(key_str: str) -> bytes:
    """Parse a key from string format. Supports hex (0x...) and ASCII."""
    if key_str.startswith("0x") or key_str.startswith("0X"):
        # Hex string
        hex_str = key_str[2:]
        if len(hex_str) % 2 != 0:
            hex_str = "0" + hex_str
        return bytes.fromhex(hex_str)
    else:
        return key_str.encode("utf-8")


def parse_hex_string(hex_str: str) -> bytes:
    """Parse a hex string (with or without 0x prefix) to bytes."""
    hex_str = hex_str.strip()
    if hex_str.startswith("0x") or hex_str.startswith("0X"):
        hex_str = hex_str[2:]
    hex_str = hex_str.replace(" ", "").replace(":", "")
    return bytes.fromhex(hex_str)


def process_input_data(input_path: str, hex_input: bool = False) -> bytes:
    """Read input data from file or stdin."""
    if input_path == "-":
        data = sys.stdin.buffer.read()
    else:
        with open(input_path, "rb") as f:
            data = f.read()

    if hex_input:
        # Input is hex-encoded text
        text = data.decode("ascii", errors="ignore").strip()
        text = text.replace(" ", "").replace("\n", "").replace("\r", "")
        data = bytes.fromhex(text)

    return data


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Decrypt obfuscated strings from malware samples",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Decryption methods:
  xor-single     Single-byte XOR (requires --key as single byte, e.g., 0x5A)
  xor-multi      Multi-byte repeating XOR (requires --key)
  xor-rolling    Rolling XOR where key evolves per byte (requires --key as initial byte)
  xor-brute      Brute-force all single-byte XOR keys, rank by printability
  base64         Standard Base64 decoding
  base64-custom  Base64 with custom alphabet (requires --alphabet)
  rc4            RC4 decryption (requires --key)
  aes-cbc        AES-CBC decryption (requires --key and --iv in hex)

Examples:
  %(prog)s --method xor-single --key 0x5A --input encrypted.bin
  %(prog)s --method xor-multi --key "secretkey" --input encrypted.bin
  %(prog)s --method xor-brute --input encrypted.bin
  %(prog)s --method base64 --input encoded_strings.txt
  %(prog)s --method rc4 --key "rc4key" --input encrypted.bin
  %(prog)s --method aes-cbc --key 0102030405060708090a0b0c0d0e0f10 --iv 00000000000000000000000000000000 --input enc.bin
        """,
    )
    parser.add_argument("--method", "-m", required=True,
                        choices=["xor-single", "xor-multi", "xor-rolling",
                                 "xor-brute", "base64", "base64-custom",
                                 "rc4", "aes-cbc"],
                        help="Decryption method to use")
    parser.add_argument("--key", "-k",
                        help="Decryption key (hex with 0x prefix or ASCII string)")
    parser.add_argument("--iv",
                        help="Initialization vector for AES-CBC (hex string)")
    parser.add_argument("--alphabet",
                        help="Custom Base64 alphabet (64 characters)")
    parser.add_argument("--input", "-i", required=True,
                        help="Input file path (use '-' for stdin)")
    parser.add_argument("--output", "-o",
                        help="Output file path (default: stdout)")
    parser.add_argument("--hex-input", action="store_true",
                        help="Treat input as hex-encoded text")
    parser.add_argument("--extract-strings", action="store_true",
                        help="Extract printable strings from decrypted output")
    parser.add_argument("--min-length", type=int, default=4,
                        help="Minimum string length for extraction (default: 4)")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")

    args = parser.parse_args()

    # Validate arguments
    if args.method in ("xor-single", "xor-multi", "xor-rolling", "rc4") and not args.key:
        parser.error(f"--key is required for method '{args.method}'")
    if args.method == "aes-cbc" and (not args.key or not args.iv):
        parser.error("--key and --iv are required for method 'aes-cbc'")
    if args.method == "base64-custom" and not args.alphabet:
        parser.error("--alphabet is required for method 'base64-custom'")

    # Read input
    try:
        data = process_input_data(args.input, args.hex_input)
    except FileNotFoundError:
        print(f"Error: Input file not found: {args.input}", file=sys.stderr)
        sys.exit(1)
    except ValueError as e:
        print(f"Error: Invalid input data: {e}", file=sys.stderr)
        sys.exit(1)

    if not data:
        print("Error: Input data is empty", file=sys.stderr)
        sys.exit(1)

    # Decrypt
    try:
        if args.method == "xor-single":
            key_bytes = parse_key(args.key)
            if len(key_bytes) != 1:
                print("Error: Single-byte XOR requires exactly one byte key (e.g., 0x5A)",
                      file=sys.stderr)
                sys.exit(1)
            result = xor_single_byte(data, key_bytes[0])

        elif args.method == "xor-multi":
            key_bytes = parse_key(args.key)
            result = xor_multi_byte(data, key_bytes)

        elif args.method == "xor-rolling":
            key_bytes = parse_key(args.key)
            initial = key_bytes[0]
            result = xor_rolling(data, initial)
            # Also try variants
            result2 = xor_rolling_variant2(data, initial)
            result3 = xor_rolling_variant3(data, initial)
            # Pick the most printable result
            candidates = [
                ("rolling (key += plaintext)", result),
                ("rolling (key += ciphertext)", result2),
                ("rolling (key ROL 1)", result3),
            ]
            best_name, best_result = max(
                candidates,
                key=lambda x: sum(1 for b in x[1] if 0x20 <= b <= 0x7E)
            )
            print(f"[*] Best rolling XOR variant: {best_name}", file=sys.stderr)
            result = best_result

        elif args.method == "xor-brute":
            print("[*] Brute-forcing single-byte XOR keys...", file=sys.stderr)
            candidates = xor_brute_force(data, args.min_length)
            if not candidates:
                print("[-] No good candidates found", file=sys.stderr)
                sys.exit(1)

            print(f"\n[+] Top candidates:", file=sys.stderr)
            for key, ratio, strings in candidates:
                print(f"\n  Key: 0x{key:02X} ({chr(key) if 0x20 <= key <= 0x7E else '?'}) "
                      f"- Printable: {ratio:.1%}", file=sys.stderr)
                for s in strings[:5]:
                    print(f"    '{s}'", file=sys.stderr)
                if len(strings) > 5:
                    print(f"    ... and {len(strings) - 5} more strings", file=sys.stderr)

            # Use best candidate for output
            best_key = candidates[0][0]
            print(f"\n[+] Using best key: 0x{best_key:02X}", file=sys.stderr)
            result = xor_single_byte(data, best_key)

        elif args.method == "base64":
            result = base64_decode(data)

        elif args.method == "base64-custom":
            result = base64_custom_decode(data, args.alphabet)

        elif args.method == "rc4":
            key_bytes = parse_key(args.key)
            result = rc4_decrypt(data, key_bytes)

        elif args.method == "aes-cbc":
            key_bytes = parse_hex_string(args.key)
            iv_bytes = parse_hex_string(args.iv)
            if len(key_bytes) not in (16, 24, 32):
                print(f"Error: AES key must be 16, 24, or 32 bytes, got {len(key_bytes)}",
                      file=sys.stderr)
                sys.exit(1)
            if len(iv_bytes) != 16:
                print(f"Error: AES IV must be 16 bytes, got {len(iv_bytes)}",
                      file=sys.stderr)
                sys.exit(1)
            result = aes_cbc_decrypt(data, key_bytes, iv_bytes)

    except Exception as e:
        print(f"Error during decryption: {e}", file=sys.stderr)
        sys.exit(1)

    # Output
    if args.extract_strings:
        strings = extract_printable_strings(result, args.min_length)
        output_text = "\n".join(strings) + "\n"
        output_data = output_text.encode("utf-8")
    else:
        output_data = result

    if args.output:
        with open(args.output, "wb") as f:
            f.write(output_data)
        print(f"[+] Output written to {args.output} ({len(output_data)} bytes)",
              file=sys.stderr)
    else:
        if args.extract_strings:
            sys.stdout.write(output_text)
        else:
            sys.stdout.buffer.write(output_data)


if __name__ == "__main__":
    main()