#!/usr/bin/env python3
"""Diff a suspicious package against its known-good version.

Recursively compares two package directories to identify added, removed, and
modified files.  For modified files, highlights suspicious changes such as
new network calls, encoded payloads, or install-hook modifications.
"""

from __future__ import annotations

import argparse
import difflib
import hashlib
import json
import re
import sys
from pathlib import Path
from typing import Any


def _sha256(path: Path) -> str:
    """Compute SHA-256 hash of a file."""
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def _relative_files(root: Path) -> set[str]:
    """Return a set of all file paths relative to *root*."""
    files: set[str] = set()
    for p in root.rglob("*"):
        if p.is_file():
            files.add(str(p.relative_to(root)))
    return files


def _is_text_file(path: Path) -> bool:
    """Heuristic check whether a file is text (not binary)."""
    try:
        chunk = path.read_bytes()[:8192]
        if b"\x00" in chunk:
            return False
        return True
    except OSError:
        return False


def _generate_diff(clean_file: Path, suspect_file: Path) -> list[str]:
    """Generate a unified diff between two text files."""
    try:
        clean_lines = clean_file.read_text(encoding="utf-8", errors="replace").splitlines(keepends=True)
        suspect_lines = suspect_file.read_text(encoding="utf-8", errors="replace").splitlines(keepends=True)
    except OSError:
        return []

    diff = list(difflib.unified_diff(
        clean_lines, suspect_lines,
        fromfile=f"clean/{clean_file.name}",
        tofile=f"suspect/{suspect_file.name}",
        lineterm="",
    ))
    return diff


def _classify_changes(diff_lines: list[str]) -> list[dict[str, str]]:
    """Classify added lines in a diff for suspicious patterns."""
    suspicious: list[dict[str, str]] = []
    patterns: dict[str, re.Pattern[str]] = {
        "dynamic_execution": re.compile(r"\b(eval|exec|compile|__import__|Function)\s*\("),
        "base64": re.compile(r"(b64decode|base64|atob)\b", re.IGNORECASE),
        "network": re.compile(r"(requests\.|urllib|http\.client|socket\.|fetch\(|XMLHttpRequest)"),
        "credential_access": re.compile(r"(os\.environ|process\.env|\.aws/|\.ssh/|\.npmrc|\.pypirc)"),
        "obfuscation": re.compile(r"(\\x[0-9a-f]{2}|chr\(\d|fromCharCode)"),
        "subprocess": re.compile(r"(subprocess|os\.system|os\.popen|child_process)"),
    }

    for line in diff_lines:
        if not line.startswith("+") or line.startswith("+++"):
            continue
        for category, pat in patterns.items():
            if pat.search(line):
                suspicious.append({
                    "category": category,
                    "line": line[1:].strip()[:150],
                })

    return suspicious


def compare_packages(
    clean_dir: Path,
    suspect_dir: Path,
) -> dict[str, Any]:
    """Compare a suspicious package directory against a clean baseline."""
    clean_files = _relative_files(clean_dir)
    suspect_files = _relative_files(suspect_dir)

    added = sorted(suspect_files - clean_files)
    removed = sorted(clean_files - suspect_files)
    common = sorted(clean_files & suspect_files)

    modified: list[dict[str, Any]] = []
    unchanged_count = 0

    for rel in common:
        clean_path = clean_dir / rel
        suspect_path = suspect_dir / rel

        clean_hash = _sha256(clean_path)
        suspect_hash = _sha256(suspect_path)

        if clean_hash == suspect_hash:
            unchanged_count += 1
            continue

        entry: dict[str, Any] = {
            "file": rel,
            "clean_hash": clean_hash,
            "suspect_hash": suspect_hash,
            "suspicious_changes": [],
            "diff_preview": [],
        }

        if _is_text_file(clean_path) and _is_text_file(suspect_path):
            diff_lines = _generate_diff(clean_path, suspect_path)
            entry["diff_preview"] = diff_lines[:50]  # first 50 lines
            entry["suspicious_changes"] = _classify_changes(diff_lines)

        modified.append(entry)

    # Risk assessment
    risk = "none"
    total_suspicious = sum(len(m["suspicious_changes"]) for m in modified)
    if added or total_suspicious >= 5:
        risk = "high"
    elif total_suspicious >= 1 or modified:
        risk = "medium"
    elif removed:
        risk = "low"

    return {
        "summary": {
            "clean_directory": str(clean_dir),
            "suspect_directory": str(suspect_dir),
            "files_added": len(added),
            "files_removed": len(removed),
            "files_modified": len(modified),
            "files_unchanged": unchanged_count,
            "risk_level": risk,
        },
        "added_files": added,
        "removed_files": removed,
        "modified_files": modified,
    }


def format_output(result: dict[str, Any], fmt: str) -> str:
    """Serialize comparison results in the requested format."""
    if fmt == "detailed":
        return json.dumps(result, indent=2)
    elif fmt == "text":
        lines: list[str] = [
            "=== Package Diff Report ===",
            f"Risk level: {result['summary']['risk_level']}",
            f"Added: {result['summary']['files_added']}",
            f"Removed: {result['summary']['files_removed']}",
            f"Modified: {result['summary']['files_modified']}",
            "",
        ]
        if result["added_files"]:
            lines.append("Added files:")
            for f in result["added_files"]:
                lines.append(f"  + {f}")
            lines.append("")
        if result["modified_files"]:
            lines.append("Modified files with suspicious changes:")
            for m in result["modified_files"]:
                if m["suspicious_changes"]:
                    lines.append(f"  * {m['file']}")
                    for c in m["suspicious_changes"]:
                        lines.append(f"    [{c['category']}] {c['line']}")
            lines.append("")
        return "\n".join(lines)
    return json.dumps(result, indent=2)


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Analyze a suspect package directory (requires separate clean dir)."""
    return {"error": "Use --clean and --suspect flags for diff analysis."}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Diff a suspicious package against a known-good version."
    )
    parser.add_argument("--input", type=Path,
                        help="Path to suspect package (alias for --suspect)")
    parser.add_argument("--clean", type=Path,
                        help="Path to known-good package directory")
    parser.add_argument("--suspect", type=Path,
                        help="Path to suspicious package directory")
    parser.add_argument("--output", type=Path,
                        help="Path to output file (stdout if omitted)")
    parser.add_argument("--format", default="json",
                        choices=["json", "text", "detailed"],
                        help="Output format (default: json)")
    args = parser.parse_args()

    suspect = args.suspect or args.input
    if not args.clean or not suspect:
        parser.error("Both --clean and --suspect (or --input) are required")

    if not args.clean.exists():
        print(f"[error] clean directory not found: {args.clean}", file=sys.stderr)
        sys.exit(1)
    if not suspect.exists():
        print(f"[error] suspect directory not found: {suspect}", file=sys.stderr)
        sys.exit(1)

    result = compare_packages(args.clean, suspect)
    rendered = format_output(result, args.format)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
        print(f"[+] Diff report written to {args.output}", file=sys.stderr)
    else:
        print(rendered)


if __name__ == "__main__":
    main()