#!/usr/bin/env python3
"""Analyze Windows driver (.sys) and Linux kernel module (.ko) files for rootkit behaviors."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import struct
import sys
from datetime import datetime
from pathlib import Path


# Suspicious strings indicating rootkit behavior
ROOTKIT_INDICATORS = {
    "hooking": [
        b"ZwQuerySystemInformation", b"NtQuerySystemInformation",
        b"ZwQueryDirectoryFile", b"NtQueryDirectoryFile",
        b"ZwEnumerateValueKey", b"NtEnumerateValueKey",
        b"ZwDeviceIoControlFile", b"KeServiceDescriptorTable",
        b"MmGetSystemRoutineAddress", b"sys_call_table",
        b"kallsyms_lookup_name", b"register_kprobe",
    ],
    "process_hiding": [
        b"PsActiveProcessHead", b"ActiveProcessLinks",
        b"EPROCESS", b"PsLookupProcessByProcessId",
        b"task_struct", b"init_task",
    ],
    "file_hiding": [
        b"IRP_MJ_DIRECTORY_CONTROL", b"FltRegisterFilter",
        b"getdents", b"filldir",
    ],
    "network_hiding": [
        b"nsiproxy", b"tcpip.sys", b"NDIS",
        b"tcp_seq_show", b"tcp4_seq_show",
    ],
    "persistence": [
        b"ZwSetValueKey", b"ZwLoadDriver",
        b"IoCreateDriver", b"DriverEntry",
        b"module_init", b"__init",
    ],
    "anti_forensics": [
        b"CmUnRegisterCallback", b"ObUnRegisterCallbacks",
        b"PsSetCreateProcessNotifyRoutine",
        b"ClearEvent", b"EventLog",
    ],
    "privilege_escalation": [
        b"SePrivilege", b"ZwOpenProcessToken",
        b"commit_creds", b"prepare_kernel_cred",
        b"current_cred",
    ],
}


def compute_hashes(filepath) -> dict:
    """Compute file hashes."""
    data = Path(filepath).read_bytes()
    return {
        "md5": hashlib.md5(data).hexdigest(),
        "sha256": hashlib.sha256(data).hexdigest(),
        "size": len(data),
    }


def extract_strings(data, min_length=4) -> dict:
    """Extract ASCII and Unicode strings."""
    strings = []

    # ASCII
    current = b""
    for byte in data:
        if 32 <= byte < 127:
            current += bytes([byte])
        else:
            if len(current) >= min_length:
                strings.append(current.decode("ascii", errors="ignore"))
            current = b""
    if len(current) >= min_length:
        strings.append(current.decode("ascii", errors="ignore"))

    return strings


def check_pe_driver(filepath) -> list:
    """Analyze a Windows PE driver file."""
    findings = []
    data = Path(filepath).read_bytes()

    # Check PE signature
    if data[:2] != b"MZ":
        return [{"type": "not_pe", "description": "File is not a PE binary"}]

    try:
        pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
        if data[pe_offset:pe_offset + 4] != b"PE\x00\x00":
            return [{"type": "invalid_pe", "description": "Invalid PE signature"}]

        # Check subsystem (native driver = 1)
        subsystem_offset = pe_offset + 0x5C
        if len(data) > subsystem_offset + 2:
            subsystem = struct.unpack_from("<H", data, subsystem_offset)[0]
            if subsystem == 1:
                findings.append({
                    "type": "native_driver",
                    "severity": "info",
                    "description": "Native subsystem (kernel driver)"
                })
    except (struct.error, IndexError):
        findings.append({
            "type": "parse_error",
            "severity": "info",
            "description": "Could not fully parse PE headers"
        })

    # Check for rootkit indicator strings
    for category, indicators in ROOTKIT_INDICATORS.items():
        for indicator in indicators:
            if indicator in data:
                findings.append({
                    "type": f"indicator_{category}",
                    "severity": "high" if category in ("hooking", "process_hiding") else "medium",
                    "description": f"Rootkit indicator ({category}): {indicator.decode('ascii', errors='ignore')}",
                    "category": category,
                    "indicator": indicator.decode("ascii", errors="ignore")
                })

    return findings


def check_elf_module(filepath) -> list:
    """Analyze a Linux ELF kernel module."""
    findings = []
    data = Path(filepath).read_bytes()

    # Check ELF magic
    if data[:4] != b"\x7fELF":
        return [{"type": "not_elf", "description": "File is not an ELF binary"}]

    findings.append({
        "type": "elf_module",
        "severity": "info",
        "description": "ELF binary detected"
    })

    # Check for rootkit indicator strings
    for category, indicators in ROOTKIT_INDICATORS.items():
        for indicator in indicators:
            if indicator in data:
                findings.append({
                    "type": f"indicator_{category}",
                    "severity": "high" if category in ("hooking", "process_hiding") else "medium",
                    "description": f"Rootkit indicator ({category}): {indicator.decode('ascii', errors='ignore')}",
                    "category": category,
                    "indicator": indicator.decode("ascii", errors="ignore")
                })

    return findings


def analyze_driver(filepath) -> dict:
    """Main driver analysis function."""
    path = Path(filepath)
    if not path.exists():
        print(f"[!] File not found: {filepath}", file=sys.stderr)
        sys.exit(1)

    data = path.read_bytes()
    hashes = compute_hashes(filepath)
    strings = extract_strings(data)

    # Determine file type and analyze
    if data[:2] == b"MZ":
        findings = check_pe_driver(filepath)
        file_type = "PE Driver (.sys)"
    elif data[:4] == b"\x7fELF":
        findings = check_elf_module(filepath)
        file_type = "ELF Module (.ko)"
    else:
        findings = [{"type": "unknown", "severity": "info", "description": "Unknown file format"}]
        file_type = "Unknown"

    # Check string entropy for obfuscation
    suspicious_strings = [s for s in strings if any(
        kw in s.lower() for kw in ["hide", "hook", "rootkit", "stealth", "inject", "shell"]
    )]

    if suspicious_strings:
        findings.append({
            "type": "suspicious_strings",
            "severity": "medium",
            "description": f"Found {len(suspicious_strings)} suspicious strings",
            "strings": suspicious_strings[:20]
        })

    return {
        "file": str(filepath),
        "file_type": file_type,
        "hashes": hashes,
        "findings": findings,
        "string_count": len(strings),
        "suspicious_string_count": len(suspicious_strings),
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Analyze driver/kernel module files for rootkit behaviors"
    )
    parser.add_argument("--input", "--file", help="Path to driver file (.sys or .ko)")
    parser.add_argument("--output", "-o", help="Output file (JSON)")
    parser.add_argument("--format", choices=["json", "text"], default="text")
    args = parser.parse_args()

    result = analyze_driver(args.file)
    result["timestamp"] = datetime.now().isoformat()

    if args.format == "json":
        output = json.dumps(result, indent=2)
    else:
        output = f"=== Driver Analysis: {result['file']} ===\n"
        output += f"Type: {result['file_type']}\n"
        output += f"SHA256: {result['hashes']['sha256']}\n"
        output += f"Size: {result['hashes']['size']} bytes\n"
        output += f"Strings: {result['string_count']} total, {result['suspicious_string_count']} suspicious\n\n"

        for f in result["findings"]:
            sev = f.get("severity", "info").upper()
            output += f"[{sev}] {f['description']}\n"

    if args.output:
        Path(args.output).write_text(output)
        print(f"[+] Report saved to {args.output}")
    else:
        print(output)


if __name__ == "__main__":
    main()