#!/usr/bin/env python3
"""
Initial triage script for suspicious file analysis.

Computes cryptographic hashes, identifies file type using magic bytes,
extracts basic metadata, and optionally queries VirusTotal API.
Outputs a structured JSON report.

Usage:
    python3 triage.py --file <path> [--output <report.json>] [--vt-lookup]
    python3 triage.py --file <path> --hashes-only
    python3 triage.py --file <path> --type-only
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import struct
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

# ---------------------------------------------------------------------------
# Optional dependency imports with graceful fallback
# ---------------------------------------------------------------------------

try:
    import magic as _magic

    def _detect_mime(path: str) -> str:
        return _magic.from_file(path, mime=True)

    def _detect_description(path: str) -> str:
        return _magic.from_file(path)

    HAS_MAGIC = True
except ImportError:
    HAS_MAGIC = False

    def _detect_mime(path: str) -> str:  # noqa: F811
        return _magic_bytes_fallback(path).get("mime_type", "application/octet-stream")

    def _detect_description(path: str) -> str:  # noqa: F811
        return _magic_bytes_fallback(path).get("description", "data")

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

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

MAGIC_SIGNATURES = {
    b"\x4d\x5a": {"description": "PE executable (MZ)", "mime_type": "application/x-dosexec"},
    b"\x7f\x45\x4c\x46": {"description": "ELF executable", "mime_type": "application/x-elf"},
    b"\xfe\xed\xfa\xce": {"description": "Mach-O executable (32-bit)", "mime_type": "application/x-mach-binary"},
    b"\xfe\xed\xfa\xcf": {"description": "Mach-O executable (64-bit)", "mime_type": "application/x-mach-binary"},
    b"\xce\xfa\xed\xfe": {"description": "Mach-O executable (32-bit, reversed)", "mime_type": "application/x-mach-binary"},
    b"\xcf\xfa\xed\xfe": {"description": "Mach-O executable (64-bit, reversed)", "mime_type": "application/x-mach-binary"},
    b"\xca\xfe\xba\xbe": {"description": "Mach-O universal binary / Java class", "mime_type": "application/x-mach-binary"},
    b"\x50\x4b\x03\x04": {"description": "ZIP archive (or Office XML / APK / JAR)", "mime_type": "application/zip"},
    b"\x50\x4b\x05\x06": {"description": "ZIP archive (empty)", "mime_type": "application/zip"},
    b"\x25\x50\x44\x46": {"description": "PDF document", "mime_type": "application/pdf"},
    b"\xd0\xcf\x11\xe0": {"description": "OLE2 compound file (legacy Office)", "mime_type": "application/x-ole-storage"},
    b"\x52\x61\x72\x21": {"description": "RAR archive", "mime_type": "application/x-rar-compressed"},
    b"\x1f\x8b": {"description": "GZIP compressed", "mime_type": "application/gzip"},
    b"\x42\x5a\x68": {"description": "BZIP2 compressed", "mime_type": "application/x-bzip2"},
    b"\x37\x7a\xbc\xaf": {"description": "7-Zip archive", "mime_type": "application/x-7z-compressed"},
    b"\x89\x50\x4e\x47": {"description": "PNG image", "mime_type": "image/png"},
    b"\xff\xd8\xff": {"description": "JPEG image", "mime_type": "image/jpeg"},
}

SUSPICIOUS_STRINGS = [
    "CreateRemoteThread",
    "VirtualAllocEx",
    "WriteProcessMemory",
    "NtUnmapViewOfSection",
    "URLDownloadToFile",
    "WinExec",
    "ShellExecute",
    "WScript.Shell",
    "powershell",
    "cmd.exe /c",
    "regsvr32",
    "rundll32",
    "mshta",
    "certutil",
    "bitsadmin",
    "net user",
    "net localgroup",
    "schtasks",
    "reg add",
    "HKEY_",
]

VT_API_BASE = "https://www.virustotal.com/api/v3"

# ---------------------------------------------------------------------------
# Hash computation
# ---------------------------------------------------------------------------

def compute_hashes(file_path: str, block_size: int = 65536) -> dict:
    """Compute MD5, SHA1, and SHA256 hashes for a file."""
    md5 = hashlib.md5()
    sha1 = hashlib.sha1()
    sha256 = hashlib.sha256()

    with open(file_path, "rb") as f:
        while True:
            block = f.read(block_size)
            if not block:
                break
            md5.update(block)
            sha1.update(block)
            sha256.update(block)

    return {
        "md5": md5.hexdigest(),
        "sha1": sha1.hexdigest(),
        "sha256": sha256.hexdigest(),
    }


def compute_ssdeep(file_path: str) -> str:
    """Compute ssdeep fuzzy hash if available."""
    try:
        import ssdeep
        return ssdeep.hash_from_file(file_path)
    except ImportError:
        return None


# ---------------------------------------------------------------------------
# File type identification (fallback without python-magic)
# ---------------------------------------------------------------------------

def _magic_bytes_fallback(file_path: str) -> dict:
    """Identify file type by reading magic bytes when python-magic is unavailable."""
    try:
        with open(file_path, "rb") as f:
            header = f.read(16)
    except OSError:
        return {"description": "unreadable", "mime_type": "application/octet-stream"}

    for sig, info in MAGIC_SIGNATURES.items():
        if header[: len(sig)] == sig:
            return info

    # Check for script types by text inspection
    try:
        text_start = header.decode("utf-8", errors="ignore").strip()
        if text_start.startswith("#!"):
            return {"description": "Script (shebang)", "mime_type": "text/x-script"}
        if text_start.startswith("<?xml") or text_start.startswith("<html"):
            return {"description": "XML/HTML document", "mime_type": "text/html"}
    except Exception:
        pass

    return {"description": "data", "mime_type": "application/octet-stream"}


def identify_file_type(file_path: str) -> dict:
    """Return file type information using python-magic or fallback."""
    if HAS_MAGIC:
        return {
            "mime_type": _detect_mime(file_path),
            "description": _detect_description(file_path),
        }
    return _magic_bytes_fallback(file_path)


# ---------------------------------------------------------------------------
# Metadata extraction
# ---------------------------------------------------------------------------

def extract_metadata(file_path: str) -> dict:
    """Extract basic file metadata."""
    stat = os.stat(file_path)
    meta = {
        "file_name": os.path.basename(file_path),
        "file_size": stat.st_size,
        "file_size_human": _human_size(stat.st_size),
        "created": datetime.fromtimestamp(stat.st_ctime, tz=timezone.utc).isoformat(),
        "modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat(),
        "platform": os.name,
    }

    # PE-specific metadata
    try:
        import pefile
        pe = pefile.PE(file_path, fast_load=True)
        pe.parse_data_directories(
            directories=[
                pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"],
                pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"],
            ]
        )
        meta["pe"] = {
            "machine": hex(pe.FILE_HEADER.Machine),
            "timestamp": datetime.fromtimestamp(
                pe.FILE_HEADER.TimeDateStamp, tz=timezone.utc
            ).isoformat(),
            "subsystem": pe.OPTIONAL_HEADER.Subsystem,
            "dll": bool(pe.FILE_HEADER.Characteristics & 0x2000),
            "num_sections": pe.FILE_HEADER.NumberOfSections,
            "entry_point": hex(pe.OPTIONAL_HEADER.AddressOfEntryPoint),
        }

        # Extract version info if present
        if hasattr(pe, "FileInfo"):
            for fi in pe.FileInfo:
                for entry in fi:
                    if hasattr(entry, "StringTable"):
                        for st in entry.StringTable:
                            for k, v in st.entries.items():
                                meta.setdefault("pe_version_info", {})[
                                    k.decode("utf-8", errors="replace")
                                ] = v.decode("utf-8", errors="replace")
        pe.close()
    except ImportError:
        pass
    except Exception:
        pass

    return meta


def _human_size(size_bytes: int) -> str:
    """Convert bytes to human-readable size."""
    for unit in ("B", "KB", "MB", "GB"):
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


# ---------------------------------------------------------------------------
# Quick string-based indicators
# ---------------------------------------------------------------------------

def quick_string_scan(file_path: str, max_bytes: int = 2 * 1024 * 1024) -> list:
    """Scan for suspicious strings in the file (limited to first max_bytes)."""
    indicators = []
    try:
        with open(file_path, "rb") as f:
            data = f.read(max_bytes)
        text = data.decode("utf-8", errors="ignore")
        text_lower = text.lower()
        for pattern in SUSPICIOUS_STRINGS:
            if pattern.lower() in text_lower:
                indicators.append(pattern)
    except OSError:
        pass
    return indicators


# ---------------------------------------------------------------------------
# VirusTotal lookup
# ---------------------------------------------------------------------------

def virustotal_lookup(sha256: str, api_key: str) -> dict:
    """Query VirusTotal API v3 for a file hash."""
    if not HAS_REQUESTS:
        return {"error": "requests library not installed; cannot query VirusTotal"}

    headers = {"x-apikey": api_key}
    url = f"{VT_API_BASE}/files/{sha256}"

    try:
        resp = requests.get(url, headers=headers, timeout=30)
        if resp.status_code == 404:
            return {"detected": False, "message": "File not found in VirusTotal"}
        if resp.status_code == 429:
            return {"error": "VirusTotal API rate limit exceeded. Try again later."}
        if resp.status_code != 200:
            return {"error": f"VirusTotal API returned status {resp.status_code}"}

        data = resp.json().get("data", {}).get("attributes", {})
        stats = data.get("last_analysis_stats", {})
        total = sum(stats.values())
        malicious = stats.get("malicious", 0) + stats.get("suspicious", 0)

        return {
            "detected": malicious > 0,
            "detections": f"{malicious}/{total}",
            "malicious": malicious,
            "undetected": stats.get("undetected", 0),
            "scan_date": data.get("last_analysis_date"),
            "popular_threat_name": data.get("popular_threat_classification", {}).get(
                "suggested_threat_label", "N/A"
            ),
            "tags": data.get("tags", []),
            "permalink": f"https://www.virustotal.com/gui/file/{sha256}",
        }
    except requests.RequestException as exc:
        return {"error": f"VirusTotal request failed: {exc}"}


# ---------------------------------------------------------------------------
# Analysis recommendation engine
# ---------------------------------------------------------------------------

def recommend_next_steps(file_info: dict, vt_result: dict | None) -> list:
    """Recommend next analysis steps based on triage findings."""
    steps = []
    mime = file_info.get("mime_type", "")
    desc = file_info.get("description", "").lower()

    # High VT detections -> minimal further analysis needed
    if vt_result and vt_result.get("malicious", 0) > 30:
        steps.append("High detection rate - extract IOCs and proceed to reporting")
        return steps

    if "executable" in desc or "dosexec" in mime or "elf" in mime:
        steps.append("Perform static analysis (PE/ELF header parsing, imports, strings)")
        steps.append("Execute in sandbox for dynamic analysis")
        steps.append("If packed/obfuscated, attempt unpacking before reverse engineering")
    elif "pdf" in mime:
        steps.append("Analyze PDF structure for JavaScript or embedded objects")
        steps.append("Extract and analyze any embedded streams")
    elif "zip" in mime or "ole" in mime:
        steps.append("Extract and inspect archive contents")
        steps.append("Check for macro content (olevba, oledump)")
        steps.append("Execute in sandbox if macros present")
    elif "script" in desc or "text" in mime:
        steps.append("Deobfuscate script content")
        steps.append("Identify C2 URLs, download URLs, or encoded payloads")
    else:
        steps.append("Determine file format with additional tools (binwalk, foremost)")
        steps.append("Perform deeper static analysis")

    if not vt_result or not vt_result.get("detected"):
        steps.append("Consider uploading to sandbox (CAPE, Any.Run, Hybrid Analysis)")

    return steps


# ---------------------------------------------------------------------------
# Main triage function
# ---------------------------------------------------------------------------

def triage(file_path: str, vt_lookup: bool = False) -> dict:
    """Run full triage on a file and return structured report."""
    if not os.path.isfile(file_path):
        return {"error": f"File not found: {file_path}"}

    report = {
        "triage_timestamp": datetime.now(tz=timezone.utc).isoformat(),
        "file_path": os.path.abspath(file_path),
    }

    # File type
    type_info = identify_file_type(file_path)
    report["file_type"] = type_info.get("description", "unknown")
    report["mime_type"] = type_info.get("mime_type", "application/octet-stream")

    # Hashes
    report["hashes"] = compute_hashes(file_path)
    ssdeep_hash = compute_ssdeep(file_path)
    if ssdeep_hash:
        report["hashes"]["ssdeep"] = ssdeep_hash

    # Metadata
    report["metadata"] = extract_metadata(file_path)
    report["file_name"] = report["metadata"]["file_name"]
    report["file_size"] = report["metadata"]["file_size"]

    # Quick indicators
    report["quick_indicators"] = quick_string_scan(file_path)

    # VirusTotal
    vt_result = None
    if vt_lookup:
        api_key = os.environ.get("VT_API_KEY", "")
        if api_key:
            vt_result = virustotal_lookup(report["hashes"]["sha256"], api_key)
            report["virustotal"] = vt_result
        else:
            report["virustotal"] = {
                "error": "VT_API_KEY environment variable not set"
            }

    # Recommendations
    report["recommended_next_steps"] = recommend_next_steps(type_info, vt_result)

    return report


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Initial triage of suspicious files",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Examples:\n"
               "  %(prog)s --file sample.exe\n"
               "  %(prog)s --file sample.exe --vt-lookup --output report.json\n"
               "  %(prog)s --file sample.exe --hashes-only\n",
    )
    parser.add_argument("--input", "--file", "-f", dest="file", required=True, help="Path to the suspicious file")
    parser.add_argument("--output", "-o", help="Write JSON report to file (default: stdout)")
    parser.add_argument("--vt-lookup", action="store_true", help="Query VirusTotal (requires VT_API_KEY env var)")
    parser.add_argument("--hashes-only", action="store_true", help="Only compute and display file hashes")
    parser.add_argument("--type-only", action="store_true", help="Only identify file type")
    parser.add_argument("--json", action="store_true", default=True, help="Output as JSON (default)")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")

    args = parser.parse_args()

    if not os.path.isfile(args.file):
        print(f"Error: File not found: {args.file}", file=sys.stderr)
        sys.exit(1)

    # -- Hashes only mode --
    if args.hashes_only:
        hashes = compute_hashes(args.file)
        result = json.dumps(hashes, indent=2)
        if args.output:
            Path(args.output).write_text(result)
        else:
            print(result)
        return

    # -- Type only mode --
    if args.type_only:
        type_info = identify_file_type(args.file)
        result = json.dumps(type_info, indent=2)
        if args.output:
            Path(args.output).write_text(result)
        else:
            print(result)
        return

    # -- Full triage --
    report = triage(args.file, vt_lookup=args.vt_lookup)
    result = json.dumps(report, indent=2, default=str)

    if args.output:
        Path(args.output).write_text(result)
        print(f"Report written to {args.output}")
    else:
        print(result)


if __name__ == "__main__":
    main()