#!/usr/bin/env python3
"""
evidence_handler.py - Hash, catalog, and maintain chain of custody for digital evidence.

Provides forensic evidence handling utilities including cryptographic hashing,
evidence cataloging, chain-of-custody logging, and integrity verification.

Usage:
    python3 evidence_handler.py --hash --input sample.exe --algorithms md5,sha1,sha256
    python3 evidence_handler.py --hash-directory --input ./evidence/ --output manifest.json
    python3 evidence_handler.py --verify --input ./evidence/ --manifest manifest.json
    python3 evidence_handler.py --init-case --case-id IR-2026-0042 --examiner "Jane Analyst"
    python3 evidence_handler.py --catalog --input disk.E01 --case-id IR-2026-0042 --evidence-id EVD-001
    python3 evidence_handler.py --log-action --case-id IR-2026-0042 --evidence-id EVD-001 --action acquired
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SUPPORTED_ALGORITHMS = ["md5", "sha1", "sha256", "sha512"]
DEFAULT_ALGORITHMS = ["md5", "sha1", "sha256"]

VALID_ACTIONS = [
    "acquired", "transferred", "analyzed", "copied",
    "stored", "returned", "destroyed", "sealed",
]


# ---------------------------------------------------------------------------
# Hashing functions
# ---------------------------------------------------------------------------

def compute_hashes(file_path: str, algorithms: Optional[list] = None) -> dict:
    """Compute cryptographic hashes of a file.

    Args:
        file_path: Path to the file to hash.
        algorithms: List of hash algorithms (md5, sha1, sha256, sha512).
                   Defaults to md5, sha1, sha256.

    Returns:
        Dictionary mapping algorithm names to hex digest strings.
    """
    if not algorithms:
        algorithms = DEFAULT_ALGORITHMS

    hashers = {}
    for algo in algorithms:
        algo = algo.lower()
        if algo not in SUPPORTED_ALGORITHMS:
            logger.warning(f"Unsupported algorithm: {algo}, skipping")
            continue
        hashers[algo] = hashlib.new(algo)

    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")

    buf_size = 65536  # 64KB chunks
    with open(path, "rb") as f:
        while True:
            data = f.read(buf_size)
            if not data:
                break
            for h in hashers.values():
                h.update(data)

    return {algo: h.hexdigest() for algo, h in hashers.items()}


def hash_directory(dir_path: str, algorithms: Optional[list] = None) -> list:
    """Compute hashes for all files in a directory.

    Args:
        dir_path: Path to the directory.
        algorithms: List of hash algorithms.

    Returns:
        List of dictionaries with file path, size, and hashes.
    """
    results = []
    base = Path(dir_path)
    if not base.is_dir():
        raise NotADirectoryError(f"Not a directory: {dir_path}")

    for filepath in sorted(base.rglob("*")):
        if filepath.is_file():
            try:
                hashes = compute_hashes(str(filepath), algorithms)
                results.append({
                    "file": str(filepath.relative_to(base)),
                    "absolute_path": str(filepath),
                    "size_bytes": filepath.stat().st_size,
                    "hashes": hashes,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                })
            except Exception as e:
                logger.error(f"Failed to hash {filepath}: {e}")
                results.append({
                    "file": str(filepath.relative_to(base)),
                    "error": str(e),
                })

    return results


def verify_manifest(dir_path: str, manifest_path: str) -> dict:
    """Verify file hashes against a previously generated manifest.

    Args:
        dir_path: Path to the evidence directory.
        manifest_path: Path to the JSON manifest file.

    Returns:
        Dictionary with verification results per file.
    """
    with open(manifest_path) as f:
        manifest = json.load(f)

    entries = manifest if isinstance(manifest, list) else manifest.get("files", [])
    results = {"verified": 0, "failed": 0, "missing": 0, "details": []}

    for entry in entries:
        if "error" in entry:
            continue

        rel_path = entry["file"]
        full_path = Path(dir_path) / rel_path

        if not full_path.exists():
            results["missing"] += 1
            results["details"].append({
                "file": rel_path,
                "status": "missing",
            })
            continue

        algorithms = list(entry["hashes"].keys())
        current_hashes = compute_hashes(str(full_path), algorithms)

        match = all(
            current_hashes.get(algo) == expected
            for algo, expected in entry["hashes"].items()
        )

        if match:
            results["verified"] += 1
            results["details"].append({"file": rel_path, "status": "verified"})
        else:
            results["failed"] += 1
            results["details"].append({
                "file": rel_path,
                "status": "FAILED",
                "expected": entry["hashes"],
                "actual": current_hashes,
            })

    results["integrity"] = "intact" if results["failed"] == 0 and results["missing"] == 0 else "COMPROMISED"
    return results


# ---------------------------------------------------------------------------
# Chain of custody
# ---------------------------------------------------------------------------

def init_case(
    case_id: str,
    examiner: str,
    description: str = "",
    output_path: Optional[str] = None,
) -> dict:
    """Initialize a new evidence case with chain-of-custody log.

    Args:
        case_id: Unique case identifier (e.g., IR-2026-0042).
        examiner: Name of the lead examiner.
        description: Case description.
        output_path: Path to write the case log JSON.

    Returns:
        Case log dictionary.
    """
    case_log = {
        "case_id": case_id,
        "created": datetime.now(timezone.utc).isoformat(),
        "examiner": examiner,
        "description": description,
        "evidence_catalog": [],
        "chain_of_custody": [],
    }

    if output_path:
        Path(output_path).write_text(json.dumps(case_log, indent=2))
        logger.info(f"Case initialized: {case_id} -> {output_path}")

    return case_log


def log_action(
    log_path: str,
    case_id: str,
    evidence_id: str,
    action: str,
    examiner: str,
    description: str = "",
    location: str = "",
    recipient: Optional[str] = None,
) -> dict:
    """Log a chain-of-custody action for an evidence item.

    Args:
        log_path: Path to the case log JSON file.
        case_id: Case identifier.
        evidence_id: Evidence item identifier.
        action: Action type (acquired, transferred, analyzed, etc.).
        examiner: Name of the person performing the action.
        description: Description of the action.
        location: Physical or logical location.
        recipient: Recipient name (for transfer actions).

    Returns:
        The new custody log entry.
    """
    log_file = Path(log_path)
    if log_file.exists():
        case_log = json.loads(log_file.read_text())
    else:
        case_log = {
            "case_id": case_id,
            "evidence_catalog": [],
            "chain_of_custody": [],
        }

    entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "evidence_id": evidence_id,
        "action": action,
        "examiner": examiner,
        "description": description,
        "location": location,
    }
    if recipient:
        entry["recipient"] = recipient

    case_log["chain_of_custody"].append(entry)
    log_file.write_text(json.dumps(case_log, indent=2))
    logger.info(f"Logged action: {action} for {evidence_id}")
    return entry


# ---------------------------------------------------------------------------
# Evidence cataloging
# ---------------------------------------------------------------------------

def catalog_evidence(
    input_path: str,
    case_id: str,
    evidence_id: str,
    description: str = "",
    examiner: str = "",
    source_path: str = "",
    tags: Optional[list] = None,
    catalog_path: Optional[str] = None,
) -> dict:
    """Catalog an evidence item with metadata and hashes.

    Args:
        input_path: Path to the evidence file.
        case_id: Case identifier.
        evidence_id: Evidence item identifier.
        description: Description of the evidence.
        examiner: Examiner who collected the evidence.
        source_path: Original path on the source system.
        tags: List of classification tags.
        catalog_path: Path to catalog JSON (appends if exists).

    Returns:
        Catalog entry dictionary.
    """
    path = Path(input_path)
    if not path.exists():
        raise FileNotFoundError(f"Evidence file not found: {input_path}")

    stat = path.stat()
    hashes = compute_hashes(input_path)

    entry = {
        "evidence_id": evidence_id,
        "case_id": case_id,
        "description": description,
        "file_name": path.name,
        "file_path": str(path.absolute()),
        "source_path": source_path,
        "size_bytes": stat.st_size,
        "hashes": hashes,
        "examiner": examiner,
        "cataloged_at": datetime.now(timezone.utc).isoformat(),
        "tags": tags or [],
    }

    if catalog_path:
        cat_file = Path(catalog_path)
        if cat_file.exists():
            catalog = json.loads(cat_file.read_text())
        else:
            catalog = {"case_id": case_id, "evidence_catalog": []}

        catalog["evidence_catalog"].append(entry)
        cat_file.write_text(json.dumps(catalog, indent=2))
        logger.info(f"Cataloged: {evidence_id} -> {catalog_path}")

    return entry


def generate_inventory(catalog_path: str, fmt: str = "json") -> str:
    """Generate an evidence inventory report from a catalog.

    Args:
        catalog_path: Path to the catalog JSON file.
        fmt: Output format (json or markdown).

    Returns:
        Formatted inventory string.
    """
    with open(catalog_path) as f:
        catalog = json.load(f)

    items = catalog.get("evidence_catalog", [])

    if fmt == "markdown":
        lines = [f"# Evidence Inventory - {catalog.get('case_id', 'Unknown')}", ""]
        lines.append("| Evidence ID | Description | SHA-256 | Size | Tags |")
        lines.append("|-------------|-------------|---------|------|------|")
        for item in items:
            sha256 = item.get("hashes", {}).get("sha256", "N/A")[:16] + "..."
            size_mb = item.get("size_bytes", 0) / (1024 * 1024)
            tags = ", ".join(item.get("tags", []))
            lines.append(
                f"| {item.get('evidence_id', 'N/A')} "
                f"| {item.get('description', 'N/A')} "
                f"| `{sha256}` "
                f"| {size_mb:.1f} MB "
                f"| {tags} |"
            )
        return "\n".join(lines)

    return json.dumps(catalog, indent=2)


def create_transfer_manifest(
    input_path: str,
    sender: str,
    recipient: str,
    case_id: str,
    output_path: Optional[str] = None,
) -> dict:
    """Create a transfer manifest with hashes for evidence exchange.

    Args:
        input_path: Path to the evidence file or package.
        sender: Name of the sender.
        recipient: Name of the recipient.
        case_id: Case identifier.
        output_path: Path to write the manifest JSON.

    Returns:
        Transfer manifest dictionary.
    """
    path = Path(input_path)
    hashes = compute_hashes(input_path)

    manifest = {
        "transfer_id": f"XFR-{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}",
        "case_id": case_id,
        "sender": sender,
        "recipient": recipient,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "file_name": path.name,
        "size_bytes": path.stat().st_size,
        "hashes": hashes,
        "status": "pending_verification",
    }

    if output_path:
        Path(output_path).write_text(json.dumps(manifest, indent=2))
        logger.info(f"Transfer manifest written to {output_path}")

    return manifest


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Digital evidence handling: hashing, cataloging, and chain-of-custody logging.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --hash --input sample.exe --algorithms md5,sha1,sha256
  %(prog)s --hash-directory --input ./evidence/ --output manifest.json
  %(prog)s --verify --input ./evidence/ --manifest manifest.json
  %(prog)s --init-case --case-id IR-2026-0042 --examiner "Jane Analyst"
  %(prog)s --catalog --input disk.E01 --case-id IR-2026-0042 --evidence-id EVD-001
  %(prog)s --log-action --case-id IR-2026-0042 --evidence-id EVD-001 --action acquired
  %(prog)s --inventory --catalog evidence_catalog.json --format markdown
        """,
    )

    # Mode selection
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--hash", action="store_true", help="Hash a single file")
    mode.add_argument("--hash-directory", action="store_true", help="Hash all files in a directory")
    mode.add_argument("--verify", action="store_true", help="Verify hashes against a manifest")
    mode.add_argument("--init-case", action="store_true", help="Initialize a new evidence case")
    mode.add_argument("--log-action", action="store_true", help="Log a chain-of-custody action")
    mode.add_argument("--catalog", action="store_true", help="Catalog an evidence item")
    mode.add_argument("--inventory", action="store_true", help="Generate evidence inventory report")
    mode.add_argument("--transfer-manifest", action="store_true", help="Create transfer manifest")
    mode.add_argument("--verify-transfer", action="store_true", help="Verify evidence against transfer manifest")

    # Common arguments
    parser.add_argument("--input", "-i", help="Input file or directory path")
    parser.add_argument("--output", "-o", help="Output file path")
    parser.add_argument("--algorithms", default="md5,sha1,sha256",
                        help="Comma-separated hash algorithms (default: md5,sha1,sha256)")

    # Case / evidence arguments
    parser.add_argument("--case-id", help="Case identifier (e.g., IR-2026-0042)")
    parser.add_argument("--evidence-id", help="Evidence item identifier (e.g., EVD-001)")
    parser.add_argument("--examiner", help="Examiner name")
    parser.add_argument("--description", help="Description text")
    parser.add_argument("--location", help="Physical/logical location")
    parser.add_argument("--action", choices=VALID_ACTIONS, help="Chain of custody action type")
    parser.add_argument("--recipient", help="Recipient name (for transfers)")
    parser.add_argument("--sender", help="Sender name (for transfers)")
    parser.add_argument("--source-path", help="Original path on source system")
    parser.add_argument("--tags", help="Comma-separated tags")
    parser.add_argument("--log", help="Path to case log JSON file")
    parser.add_argument("--manifest", help="Path to hash manifest for verification")
    parser.add_argument("--format", choices=["json", "markdown"], default="json",
                        help="Output format (default: json)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")

    return parser.parse_args()


def main() -> None:
    args = parse_arguments()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    if args.hash:
        if not args.input:
            logger.error("--input required for --hash")
            sys.exit(1)
        algorithms = [a.strip() for a in args.algorithms.split(",")]
        result = compute_hashes(args.input, algorithms)
        output = {
            "file": args.input,
            "size_bytes": Path(args.input).stat().st_size,
            "hashes": result,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        output_json = json.dumps(output, indent=2)
        if args.output:
            Path(args.output).write_text(output_json)
            logger.info(f"Hashes written to {args.output}")
        else:
            print(output_json)

    elif args.hash_directory:
        if not args.input:
            logger.error("--input required for --hash-directory")
            sys.exit(1)
        algorithms = [a.strip() for a in args.algorithms.split(",")]
        results = hash_directory(args.input, algorithms)
        output = {"directory": args.input, "files": results, "total": len(results)}
        output_json = json.dumps(output, indent=2)
        if args.output:
            Path(args.output).write_text(output_json)
            logger.info(f"Directory manifest written to {args.output}")
        else:
            print(output_json)

    elif args.verify:
        if not args.input or not args.manifest:
            logger.error("--input and --manifest required for --verify")
            sys.exit(1)
        result = verify_manifest(args.input, args.manifest)
        print(json.dumps(result, indent=2))
        sys.exit(0 if result["integrity"] == "intact" else 1)

    elif args.init_case:
        if not args.case_id or not args.examiner:
            logger.error("--case-id and --examiner required for --init-case")
            sys.exit(1)
        result = init_case(
            case_id=args.case_id,
            examiner=args.examiner,
            description=args.description or "",
            output_path=args.output,
        )
        if not args.output:
            print(json.dumps(result, indent=2))

    elif args.log_action:
        if not all([args.case_id, args.evidence_id, args.action, args.examiner, args.log]):
            logger.error("--case-id, --evidence-id, --action, --examiner, and --log required")
            sys.exit(1)
        entry = log_action(
            log_path=args.log,
            case_id=args.case_id,
            evidence_id=args.evidence_id,
            action=args.action,
            examiner=args.examiner,
            description=args.description or "",
            location=args.location or "",
            recipient=args.recipient,
        )
        print(json.dumps(entry, indent=2))

    elif args.catalog:
        if not all([args.input, args.case_id, args.evidence_id]):
            logger.error("--input, --case-id, and --evidence-id required for --catalog")
            sys.exit(1)
        tags = [t.strip() for t in args.tags.split(",")] if args.tags else []
        entry = catalog_evidence(
            input_path=args.input,
            case_id=args.case_id,
            evidence_id=args.evidence_id,
            description=args.description or "",
            examiner=args.examiner or "",
            source_path=args.source_path or "",
            tags=tags,
            catalog_path=args.output,
        )
        print(json.dumps(entry, indent=2))

    elif args.inventory:
        catalog_file = args.output or args.input
        if not catalog_file:
            logger.error("Provide catalog path via --input or --catalog flag")
            sys.exit(1)
        # Use a separate --catalog arg or fallback
        cat_path = args.manifest or args.input
        if not cat_path:
            logger.error("Provide catalog path")
            sys.exit(1)
        report = generate_inventory(cat_path, args.format)
        print(report)

    elif args.transfer_manifest:
        if not all([args.input, args.sender, args.recipient, args.case_id]):
            logger.error("--input, --sender, --recipient, and --case-id required")
            sys.exit(1)
        manifest = create_transfer_manifest(
            input_path=args.input,
            sender=args.sender,
            recipient=args.recipient,
            case_id=args.case_id,
            output_path=args.output,
        )
        if not args.output:
            print(json.dumps(manifest, indent=2))

    elif args.verify_transfer:
        if not args.input or not args.manifest:
            logger.error("--input and --manifest required for --verify-transfer")
            sys.exit(1)
        with open(args.manifest) as f:
            manifest = json.load(f)
        current_hashes = compute_hashes(args.input)
        match = all(
            current_hashes.get(algo) == expected
            for algo, expected in manifest.get("hashes", {}).items()
        )
        result = {
            "file": args.input,
            "transfer_id": manifest.get("transfer_id"),
            "integrity": "verified" if match else "FAILED",
            "expected_hashes": manifest.get("hashes"),
            "actual_hashes": current_hashes,
        }
        print(json.dumps(result, indent=2))
        sys.exit(0 if match else 1)

    else:
        logger.error("Specify a mode: --hash, --hash-directory, --verify, --init-case, "
                      "--log-action, --catalog, --inventory, --transfer-manifest, --verify-transfer")
        sys.exit(1)


if __name__ == "__main__":
    main()