#!/usr/bin/env python3
"""Test YARA rules against malware samples and clean file corpora.

Validates detection accuracy by scanning known malware directories for true
positives and clean file directories for false positives, producing a
structured test report.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any


def find_files(directory: Path, recursive: bool = True) -> list[Path]:
    """Find all files in a directory, optionally recursively."""
    if not directory.exists():
        return []
    if recursive:
        return [p for p in directory.rglob("*") if p.is_file()]
    return [p for p in directory.iterdir() if p.is_file()]


def load_yara_rules(rules_path: Path) -> Any:
    """Load YARA rules from a file.

    Attempts to use yara-python if available, otherwise falls back to
    a basic syntax validation stub.
    """
    try:
        import yara  # type: ignore[import-untyped]
        return yara.compile(filepath=str(rules_path))
    except ImportError:
        # Fallback: validate the rule file is readable
        rules_text = rules_path.read_text(encoding="utf-8")
        if "rule " not in rules_text:
            raise ValueError(f"No YARA rules found in {rules_path}")
        return None
    except Exception as e:
        print(f"Error compiling YARA rules: {e}", file=sys.stderr)
        raise


def scan_file(rules: Any, file_path: Path) -> list[dict[str, str]]:
    """Scan a single file with YARA rules. Returns list of matches."""
    if rules is None:
        # Stub mode: no yara-python available
        return []
    try:
        matches = rules.match(filepath=str(file_path))
        return [
            {
                "rule": str(match.rule),
                "tags": list(match.tags) if match.tags else [],
                "namespace": str(match.namespace),
            }
            for match in matches
        ]
    except Exception as e:
        return [{"error": str(e)}]


def test_rules(
    rules_path: Path,
    malware_dir: Path | None = None,
    clean_dir: Path | None = None,
) -> dict[str, Any]:
    """Test YARA rules against malware and clean file directories."""
    start_time = time.time()

    rules = load_yara_rules(rules_path)
    yara_available = rules is not None

    result: dict[str, Any] = {
        "rules_file": str(rules_path),
        "yara_available": yara_available,
        "true_positives": [],
        "false_negatives": [],
        "false_positives": [],
        "true_negatives": 0,
        "statistics": {},
    }

    # Test against known malware samples
    malware_files: list[Path] = []
    if malware_dir and malware_dir.exists():
        malware_files = find_files(malware_dir)
        for mf in malware_files:
            matches = scan_file(rules, mf)
            if matches and not any("error" in m for m in matches):
                result["true_positives"].append({
                    "file": str(mf),
                    "matches": matches,
                })
            else:
                result["false_negatives"].append({
                    "file": str(mf),
                    "note": "No rule matched this known malware sample",
                })

    # Test against clean files
    clean_files: list[Path] = []
    if clean_dir and clean_dir.exists():
        clean_files = find_files(clean_dir)
        for cf in clean_files:
            matches = scan_file(rules, cf)
            if matches and not any("error" in m for m in matches):
                result["false_positives"].append({
                    "file": str(cf),
                    "matches": matches,
                })
            else:
                result["true_negatives"] += 1

    elapsed = time.time() - start_time

    tp = len(result["true_positives"])
    fn = len(result["false_negatives"])
    fp = len(result["false_positives"])
    tn = result["true_negatives"]

    detection_rate = (tp / (tp + fn) * 100) if (tp + fn) > 0 else 0.0
    fp_rate = (fp / (fp + tn) * 100) if (fp + tn) > 0 else 0.0
    precision = (tp / (tp + fp) * 100) if (tp + fp) > 0 else 0.0

    result["statistics"] = {
        "malware_files_scanned": len(malware_files),
        "clean_files_scanned": len(clean_files),
        "true_positives": tp,
        "false_negatives": fn,
        "false_positives": fp,
        "true_negatives": tn,
        "detection_rate_pct": round(detection_rate, 2),
        "false_positive_rate_pct": round(fp_rate, 2),
        "precision_pct": round(precision, 2),
        "elapsed_seconds": round(elapsed, 3),
    }

    if not yara_available:
        result["warning"] = (
            "yara-python is not installed. Install it with: pip install yara-python. "
            "Rule syntax was validated but no file scanning was performed."
        )

    return result


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Analyze YARA rules from the given file."""
    return test_rules(input_path)


def main() -> None:
    """Entry point for the YARA rule tester CLI."""
    parser = argparse.ArgumentParser(
        description="Test YARA rules against malware samples and clean files."
    )
    parser.add_argument(
        "--input",
        "--rules",
        type=Path,
        required=True,
        dest="input",
        help="Path to YARA rules 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)",
    )
    parser.add_argument(
        "--malware-dir",
        type=Path,
        help="Directory containing known malware samples",
    )
    parser.add_argument(
        "--clean-dir",
        type=Path,
        help="Directory containing clean/benign files for FP testing",
    )
    args = parser.parse_args()

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

    result = test_rules(
        rules_path=args.input,
        malware_dir=args.malware_dir,
        clean_dir=args.clean_dir,
    )

    if args.format == "text":
        output_text = _format_text(result)
    elif args.format == "csv":
        output_text = _format_csv(result)
    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"Results written to {args.output}", file=sys.stderr)
    else:
        print(output_text)


def _format_text(result: dict[str, Any]) -> str:
    """Format test results as human-readable text."""
    stats = result["statistics"]
    lines: list[str] = [
        "=== YARA Rule Test Results ===",
        f"Rules file: {result['rules_file']}",
        f"YARA available: {result['yara_available']}",
        "",
        f"Malware files scanned: {stats['malware_files_scanned']}",
        f"Clean files scanned:   {stats['clean_files_scanned']}",
        "",
        f"True positives:    {stats['true_positives']}",
        f"False negatives:   {stats['false_negatives']}",
        f"False positives:   {stats['false_positives']}",
        f"True negatives:    {stats['true_negatives']}",
        "",
        f"Detection rate:    {stats['detection_rate_pct']}%",
        f"False positive rate: {stats['false_positive_rate_pct']}%",
        f"Precision:         {stats['precision_pct']}%",
        f"Elapsed:           {stats['elapsed_seconds']}s",
    ]

    if result.get("warning"):
        lines.extend(["", f"WARNING: {result['warning']}"])

    if result["false_positives"]:
        lines.extend(["", "False Positive Files:"])
        for fp in result["false_positives"]:
            lines.append(f"  - {fp['file']}")

    if result["false_negatives"]:
        lines.extend(["", "Missed Malware Files:"])
        for fn_item in result["false_negatives"]:
            lines.append(f"  - {fn_item['file']}")

    return "\n".join(lines)


def _format_csv(result: dict[str, Any]) -> str:
    """Format test results as CSV."""
    lines: list[str] = ["type,file,rule,detail"]
    for tp in result["true_positives"]:
        for m in tp["matches"]:
            lines.append(f"TP,{tp['file']},{m.get('rule', 'N/A')},detected")
    for fn_item in result["false_negatives"]:
        lines.append(f"FN,{fn_item['file']},none,missed")
    for fp in result["false_positives"]:
        for m in fp["matches"]:
            lines.append(f"FP,{fp['file']},{m.get('rule', 'N/A')},false_alarm")
    return "\n".join(lines)


if __name__ == "__main__":
    main()