#!/usr/bin/env python3
"""Generate YARA rules from malware samples by extracting unique strings and byte patterns.

Analyzes a binary sample to identify candidate detection strings, byte sequences,
and structural characteristics, then generates a YARA rule template with appropriate
conditions and metadata.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# Strings commonly found in legitimate software (high false-positive risk)
COMMON_STRINGS: set[str] = {
    "createfile",
    "regsetvalue",
    "error",
    "failed",
    "success",
    "http/1.1",
    "content-type",
    "user-agent",
    "accept",
    "microsoft",
    "windows",
    "kernel32.dll",
    "ntdll.dll",
    "advapi32.dll",
    "copyright",
    "version",
}

# PE magic bytes
PE_MAGIC = b"MZ"


def compute_hashes(data: bytes) -> dict[str, str]:
    """Compute MD5, SHA1, and SHA256 hashes of binary data."""
    return {
        "md5": hashlib.md5(data).hexdigest(),
        "sha1": hashlib.sha1(data).hexdigest(),
        "sha256": hashlib.sha256(data).hexdigest(),
    }


def extract_ascii_strings(data: bytes, min_length: int = 6) -> list[str]:
    """Extract printable ASCII strings of at least min_length characters."""
    pattern = rb"[\x20-\x7E]{" + str(min_length).encode() + rb",}"
    return [m.decode("ascii", errors="ignore") for m in re.findall(pattern, data)]


def extract_unicode_strings(data: bytes, min_length: int = 6) -> list[str]:
    """Extract UTF-16LE encoded strings of at least min_length characters."""
    pattern = rb"(?:[\x20-\x7E]\x00){" + str(min_length).encode() + rb",}"
    matches = re.findall(pattern, data)
    return [m.decode("utf-16-le", errors="ignore") for m in matches]


def score_string(s: str) -> float:
    """Score a string for its suitability as a YARA detection string.

    Higher scores indicate more unique/valuable strings.
    """
    score = 0.0

    # Length bonus
    if len(s) >= 10:
        score += 2.0
    elif len(s) >= 7:
        score += 1.0

    # Penalize common strings
    if s.lower() in COMMON_STRINGS:
        score -= 10.0

    # Bonus for likely malware indicators
    malware_indicators = [
        "mutex", "c2", "beacon", "shell", "inject", "payload",
        "ransom", "crypt", "exfil", "keylog", "rat", "bot",
        "cmd_", "download", "upload", "execute", "persist",
        ".onion", "tor2web",
    ]
    for indicator in malware_indicators:
        if indicator in s.lower():
            score += 5.0
            break

    # Bonus for paths (PDB, file paths)
    if "\\" in s and ("Users" in s or "Desktop" in s or ".pdb" in s.lower()):
        score += 4.0

    # Bonus for URLs/domains
    if re.search(r"https?://|[a-z0-9\-]+\.[a-z]{2,}", s, re.IGNORECASE):
        score += 3.0

    # Bonus for mutex-like patterns
    if s.startswith("Global\\") or s.startswith("Local\\"):
        score += 5.0

    # Penalize very short strings
    if len(s) < 6:
        score -= 5.0

    return score


def select_candidate_strings(
    ascii_strings: list[str],
    unicode_strings: list[str],
    max_strings: int = 15,
) -> list[dict[str, Any]]:
    """Select the best candidate strings for YARA rule generation."""
    candidates: list[dict[str, Any]] = []

    for s in ascii_strings:
        candidates.append({
            "value": s,
            "encoding": "ascii",
            "score": score_string(s),
        })

    for s in unicode_strings:
        candidates.append({
            "value": s,
            "encoding": "wide",
            "score": score_string(s),
        })

    # Sort by score descending and take top candidates
    candidates.sort(key=lambda x: x["score"], reverse=True)
    return [c for c in candidates[:max_strings] if c["score"] > 0]


def is_pe_file(data: bytes) -> bool:
    """Check if binary data starts with PE magic bytes."""
    return data[:2] == PE_MAGIC


def generate_rule(
    sample_path: Path,
    rule_name: str | None = None,
    author: str = "Auto-generated",
    max_strings: int = 15,
) -> dict[str, Any]:
    """Generate a YARA rule from the given sample."""
    data = sample_path.read_bytes()
    hashes = compute_hashes(data)

    ascii_strings = extract_ascii_strings(data)
    unicode_strings = extract_unicode_strings(data)
    candidates = select_candidate_strings(ascii_strings, unicode_strings, max_strings)

    if rule_name is None:
        safe_name = re.sub(r"[^a-zA-Z0-9_]", "_", sample_path.stem)
        rule_name = f"mal_{safe_name}"

    # Build YARA rule text
    rule_lines: list[str] = []
    rule_lines.append(f"rule {rule_name} {{")

    # Metadata
    rule_lines.append("    meta:")
    rule_lines.append(f'        author = "{author}"')
    rule_lines.append(f'        date = "{datetime.now(timezone.utc).strftime("%Y-%m-%d")}"')
    rule_lines.append(f'        description = "Auto-generated rule for {sample_path.name}"')
    rule_lines.append(f'        hash = "{hashes["sha256"]}"')
    rule_lines.append('        tlp = "WHITE"')

    # Strings
    if candidates:
        rule_lines.append("")
        rule_lines.append("    strings:")
        for i, candidate in enumerate(candidates):
            var_name = f"$s{i}"
            value = candidate["value"].replace("\\", "\\\\").replace('"', '\\"')
            modifier = "ascii wide" if candidate["encoding"] == "wide" else "ascii"
            rule_lines.append(f'        {var_name} = "{value}" {modifier}')

    # Condition
    rule_lines.append("")
    rule_lines.append("    condition:")

    conditions: list[str] = []
    if is_pe_file(data):
        conditions.append("uint16(0) == 0x5A4D")

    # File size condition (within 2x of sample size)
    size_mb = len(data) / (1024 * 1024)
    max_size = max(int(size_mb * 2) + 1, 1)
    conditions.append(f"filesize < {max_size}MB")

    if candidates:
        min_matches = max(2, len(candidates) // 3)
        conditions.append(f"{min_matches} of ($s*)")

    rule_lines.append("        " + " and\n        ".join(conditions))
    rule_lines.append("}")

    rule_text = "\n".join(rule_lines)

    return {
        "rule_name": rule_name,
        "rule_text": rule_text,
        "sample": str(sample_path),
        "hashes": hashes,
        "candidate_strings": len(candidates),
        "is_pe": is_pe_file(data),
        "file_size": len(data),
    }


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Analyze a sample and return the generated YARA rule."""
    return generate_rule(input_path)


def main() -> None:
    """Entry point for the YARA generator CLI."""
    parser = argparse.ArgumentParser(
        description="Generate YARA rules from malware samples."
    )
    parser.add_argument(
        "--input",
        "--sample",
        type=Path,
        required=True,
        dest="input",
        help="Path to the malware sample",
    )
    parser.add_argument(
        "--output", type=Path, help="Path to output file (stdout if omitted)"
    )
    parser.add_argument(
        "--format",
        default="json",
        choices=["json", "yara", "text"],
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--rule-name",
        type=str,
        help="Custom rule name (auto-generated if omitted)",
    )
    parser.add_argument(
        "--author",
        type=str,
        default="Auto-generated",
        help="Rule author name",
    )
    parser.add_argument(
        "--max-strings",
        type=int,
        default=15,
        help="Maximum number of strings to include (default: 15)",
    )
    args = parser.parse_args()

    if not args.input.exists():
        print(f"Error: Sample not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    result = generate_rule(
        args.input,
        rule_name=args.rule_name,
        author=args.author,
        max_strings=args.max_strings,
    )

    if args.format == "yara":
        output_text = result["rule_text"]
    elif args.format == "text":
        output_text = (
            f"Generated YARA Rule: {result['rule_name']}\n"
            f"Sample: {result['sample']}\n"
            f"SHA256: {result['hashes']['sha256']}\n"
            f"Strings selected: {result['candidate_strings']}\n"
            f"PE file: {result['is_pe']}\n\n"
            f"{result['rule_text']}\n"
        )
    else:
        output_text = json.dumps(result, indent=2)

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


if __name__ == "__main__":
    main()