#!/usr/bin/env python3
"""String decryptor — decrypt and deobfuscate strings in malware binaries."""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import math
import re
import struct
import sys
from collections import Counter
from pathlib import Path


def calculate_entropy(data: bytes) -> float:
    """Calculate Shannon entropy of byte data."""
    if not data:
        return 0.0
    freq = Counter(data)
    length = len(data)
    return -sum((c / length) * math.log2(c / length) for c in freq.values())


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


def xor_multi_byte(data: bytes, key: bytes) -> bytes:
    """XOR decrypt with a multi-byte key."""
    return bytes(b ^ key[i % len(key)] for i, b in enumerate(data))


def xor_rolling(data: bytes, initial_key: int) -> bytes:
    """XOR decrypt with a rolling key (each byte XORed with previous ciphertext)."""
    result = bytearray()
    key = initial_key
    for b in data:
        decrypted = b ^ key
        result.append(decrypted)
        key = b  # Key becomes previous ciphertext byte
    return bytes(result)


def rc4_decrypt(data: bytes, key: bytes) -> bytes:
    """RC4 stream cipher decryption."""
    S = list(range(256))
    j = 0
    for i in range(256):
        j = (j + S[i] + key[i % len(key)]) % 256
        S[i], S[j] = S[j], S[i]

    i = j = 0
    result = bytearray()
    for byte in data:
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i]
        k = S[(S[i] + S[j]) % 256]
        result.append(byte ^ k)
    return bytes(result)


def decode_base64_strings(data: bytes) -> list[dict]:
    """Find and decode base64-encoded strings in binary data."""
    results = []
    b64_pattern = rb'[A-Za-z0-9+/]{16,}={0,2}'
    for match in re.finditer(b64_pattern, data):
        try:
            decoded = base64.b64decode(match.group())
            text = decoded.decode("utf-8", errors="strict")
            if text.isprintable() and len(text) >= 4:
                results.append({
                    "offset": match.start(),
                    "encoded": match.group().decode("ascii"),
                    "decoded": text,
                    "encoding": "base64",
                })
        except Exception:
            pass
    return results


def extract_stack_strings(data: bytes) -> list[dict]:
    """Detect potential stack string construction patterns.

    Stack strings are built character-by-character using mov instructions.
    Pattern: mov [ebp-XX], byte_value repeated sequences.
    """
    results = []
    # Look for sequences of byte pushes/moves that could be stack strings
    # Pattern: C6 45 XX YY = mov byte [ebp+XX], YY (x86)
    pattern = rb'(\xc6\x45.[\x20-\x7e]){4,}'
    for match in re.finditer(pattern, data):
        chars = []
        pos = match.start()
        while pos + 3 < match.end():
            if data[pos] == 0xC6 and data[pos + 1] == 0x45:
                chars.append(chr(data[pos + 3]))
                pos += 4
            else:
                break
        if len(chars) >= 4:
            results.append({
                "offset": match.start(),
                "string": "".join(chars),
                "type": "stack_string",
                "length": len(chars),
            })
    return results


def brute_force_xor(data: bytes, min_length: int = 6) -> list[dict]:
    """Try all single-byte XOR keys and look for readable strings."""
    results = []
    for key in range(1, 256):
        decrypted = xor_single_byte(data, key)
        # Look for ASCII string runs
        current_string = []
        for i, b in enumerate(decrypted):
            if 0x20 <= b < 0x7F:
                current_string.append(chr(b))
            else:
                if len(current_string) >= min_length:
                    s = "".join(current_string)
                    # Filter for interesting strings
                    if any(kw in s.lower() for kw in ("http", "www", ".com", ".exe", "cmd",
                                                       "password", "admin", "user", "key")):
                        results.append({
                            "xor_key": f"0x{key:02x}",
                            "offset": i - len(current_string),
                            "string": s,
                        })
                current_string = []
    return results


def find_encrypted_blobs(data: bytes, min_size: int = 32) -> list[dict]:
    """Find high-entropy regions that may contain encrypted data."""
    blobs = []
    window_size = 256
    threshold = 7.0

    for offset in range(0, len(data) - window_size, window_size // 2):
        window = data[offset:offset + window_size]
        entropy = calculate_entropy(window)
        if entropy >= threshold:
            blobs.append({
                "offset": offset,
                "size": window_size,
                "entropy": round(entropy, 2),
            })

    # Merge adjacent blobs
    merged = []
    for blob in blobs:
        if merged and blob["offset"] <= merged[-1]["offset"] + merged[-1]["size"]:
            merged[-1]["size"] = blob["offset"] + blob["size"] - merged[-1]["offset"]
            merged[-1]["entropy"] = max(merged[-1]["entropy"], blob["entropy"])
        else:
            merged.append(blob)

    return [b for b in merged if b["size"] >= min_size]


def main() -> None:
    parser = argparse.ArgumentParser(description="Malware String Decryptor")
    parser.add_argument("--input", required=True, help="Input binary file")
    parser.add_argument(
        "--mode",
        choices=["xor-brute", "base64", "stack-strings", "rc4", "blobs", "all"],
        default="all",
        help="Decryption mode",
    )
    parser.add_argument("--key", help="Decryption key (hex string, e.g., 'deadbeef')")
    parser.add_argument("--key-type", choices=["xor", "rc4"], default="xor")
    parser.add_argument("--min-length", type=int, default=6, help="Minimum string length")
    parser.add_argument("--output", default="decrypted_strings.json", help="Output file path")
    parser.add_argument("--format", choices=["json", "csv", "text"], default="json")

    args = parser.parse_args()

    input_path = Path(args.input)
    if not input_path.exists():
        print(f"[!] File not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    data = input_path.read_bytes()
    sha256 = hashlib.sha256(data).hexdigest()

    print(f"[*] String Decryptor — mode: {args.mode}")
    print(f"[*] Input: {args.input} ({len(data)} bytes)")
    print(f"[*] SHA-256: {sha256}")

    results = {
        "input_file": str(input_path),
        "sha256": sha256,
        "size_bytes": len(data),
        "mode": args.mode,
        "decrypted_strings": [],
    }

    if args.key:
        key_bytes = bytes.fromhex(args.key)
        if args.key_type == "xor":
            decrypted = xor_multi_byte(data, key_bytes)
        elif args.key_type == "rc4":
            decrypted = rc4_decrypt(data, key_bytes)
        # Extract printable strings from decrypted data
        for match in re.finditer(rb'[\x20-\x7e]{6,}', decrypted):
            results["decrypted_strings"].append({
                "offset": match.start(),
                "string": match.group().decode("ascii"),
                "method": args.key_type,
                "key": args.key,
            })
    else:
        if args.mode in ("xor-brute", "all"):
            print("[*] Brute-forcing single-byte XOR keys...")
            results["xor_results"] = brute_force_xor(data, args.min_length)
            print(f"    Found {len(results['xor_results'])} interesting strings")

        if args.mode in ("base64", "all"):
            print("[*] Scanning for base64-encoded strings...")
            results["base64_results"] = decode_base64_strings(data)
            print(f"    Found {len(results['base64_results'])} base64 strings")

        if args.mode in ("stack-strings", "all"):
            print("[*] Detecting stack strings...")
            results["stack_strings"] = extract_stack_strings(data)
            print(f"    Found {len(results['stack_strings'])} stack strings")

        if args.mode in ("blobs", "all"):
            print("[*] Finding encrypted blobs...")
            results["encrypted_blobs"] = find_encrypted_blobs(data)
            print(f"    Found {len(results['encrypted_blobs'])} high-entropy regions")

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if args.format == "text":
        with open(output_path, "w") as f:
            for key in ("xor_results", "base64_results", "stack_strings", "decrypted_strings"):
                for item in results.get(key, []):
                    f.write(f"{item.get('string', item.get('decoded', ''))}\n")
    else:
        output_path.write_text(json.dumps(results, indent=2))

    print(f"[*] Results written to {args.output}")


if __name__ == "__main__":
    main()