#!/usr/bin/env python3
"""Extract and analyze firmware images.

Wraps binwalk for firmware extraction and performs security-focused
analysis: finding hardcoded credentials, suspicious binaries, and
network services.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path


def compute_hash(filepath) -> dict:
    """Compute SHA256 hash of a file."""
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def run_binwalk(firmware_path, output_dir) -> dict:
    """Run binwalk to extract firmware."""
    results = {"binwalk_available": False, "signatures": [], "extracted": False}

    try:
        # Signature scan
        sig_result = subprocess.run(
            ["binwalk", str(firmware_path)],
            capture_output=True, text=True, timeout=120
        )
        results["binwalk_available"] = True
        results["signatures_raw"] = sig_result.stdout

        # Parse signatures
        for line in sig_result.stdout.strip().split("\n")[3:]:
            parts = line.strip().split(None, 2)
            if len(parts) >= 3:
                results["signatures"].append({
                    "offset": parts[0],
                    "type": parts[2] if len(parts) > 2 else parts[1],
                })

        # Extract
        ext_result = subprocess.run(
            ["binwalk", "-e", "-C", str(output_dir), str(firmware_path)],
            capture_output=True, text=True, timeout=300
        )
        results["extracted"] = ext_result.returncode == 0

        # Identify architecture
        arch_result = subprocess.run(
            ["binwalk", "-A", str(firmware_path)],
            capture_output=True, text=True, timeout=60
        )
        results["architecture_scan"] = arch_result.stdout

    except FileNotFoundError:
        results["error"] = "binwalk not installed. Install: sudo apt install binwalk"
    except subprocess.TimeoutExpired:
        results["error"] = "binwalk timed out"

    return results


def find_credentials(extracted_dir) -> list:
    """Search extracted filesystem for hardcoded credentials."""
    findings = []
    search_patterns = [
        (r"root:.*:\d+:\d+:", "Root account in passwd"),
        (r"admin:.*:\d+:\d+:", "Admin account in passwd"),
        (r"\$[156]\$[a-zA-Z0-9./]+\$", "Password hash found"),
        (r"password\s*[=:]\s*\S+", "Hardcoded password"),
        (r"api[_-]?key\s*[=:]\s*\S+", "API key"),
        (r"secret\s*[=:]\s*\S+", "Secret value"),
    ]

    for root, dirs, files in os.walk(extracted_dir):
        for fname in files:
            filepath = os.path.join(root, fname)
            try:
                with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
                    content = f.read(50000)  # Limit read size

                for pattern, description in search_patterns:
                    matches = re.findall(pattern, content, re.IGNORECASE)
                    if matches:
                        findings.append({
                            "file": filepath,
                            "type": description,
                            "matches": matches[:5],
                            "severity": "critical" if "password" in description.lower() else "high",
                        })
            except (OSError, UnicodeDecodeError):
                continue

    return findings


def find_network_services(extracted_dir) -> list:
    """Identify network services in extracted firmware."""
    findings = []
    services = [
        "telnetd", "sshd", "httpd", "lighttpd", "nginx",
        "ftpd", "tftpd", "snmpd", "upnpd", "miniupnpd",
    ]

    for root, dirs, files in os.walk(extracted_dir):
        for fname in files:
            filepath = os.path.join(root, fname)
            if fname in services or any(s in fname for s in services):
                findings.append({
                    "file": filepath,
                    "service": fname,
                    "severity": "medium",
                    "description": f"Network service binary: {fname}",
                })

    # Check init scripts
    init_dirs = ["etc/init.d", "etc/rc.d", "etc/default"]
    for init_dir in init_dirs:
        full_path = os.path.join(extracted_dir, init_dir)
        if os.path.exists(full_path):
            for fname in os.listdir(full_path):
                filepath = os.path.join(full_path, fname)
                try:
                    content = Path(filepath).read_text(errors="ignore")
                    for svc in services:
                        if svc in content:
                            findings.append({
                                "file": filepath,
                                "service": svc,
                                "severity": "info",
                                "description": f"Init script references {svc}",
                            })
                except OSError:
                    continue

    return findings


def find_suspicious_binaries(extracted_dir) -> list:
    """Find executables and analyze with 'file' command."""
    findings = []

    for root, dirs, files in os.walk(extracted_dir):
        for fname in files:
            filepath = os.path.join(root, fname)
            try:
                # Check if executable
                if os.access(filepath, os.X_OK) or fname.endswith((".elf", ".bin")):
                    result = subprocess.run(
                        ["file", filepath],
                        capture_output=True, text=True, timeout=5
                    )
                    file_type = result.stdout.strip()

                    if "ELF" in file_type or "executable" in file_type:
                        findings.append({
                            "file": filepath,
                            "type": file_type.split(":", 1)[-1].strip()[:200],
                            "size": os.path.getsize(filepath),
                            "hash": compute_hash(filepath),
                        })
            except (OSError, subprocess.TimeoutExpired):
                continue

    return findings


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract and analyze firmware images"
    )
    parser.add_argument("--input", "--image", "-i", required=True, help="Firmware image file")
    parser.add_argument("--output", "-o", default="./extracted", help="Output directory")
    parser.add_argument("--report", "-r", help="Report output file (JSON)")
    parser.add_argument("--format", choices=["json", "text"], default="text")
    args = parser.parse_args()

    image_path = Path(args.image)
    if not image_path.exists():
        print(f"[!] File not found: {args.image}", file=sys.stderr)
        sys.exit(1)

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    report = {
        "tool": "firmware_extractor",
        "timestamp": datetime.now().isoformat(),
        "firmware_file": str(args.image),
        "firmware_hash": compute_hash(args.image),
        "firmware_size": image_path.stat().st_size,
    }

    # Run binwalk extraction
    print("[*] Running binwalk extraction...")
    report["binwalk"] = run_binwalk(args.image, output_dir)

    # Find extracted directory
    extracted_dirs = list(output_dir.rglob("*squashfs*")) + list(output_dir.rglob("*filesystem*"))
    if not extracted_dirs:
        extracted_dirs = [d for d in output_dir.iterdir() if d.is_dir()]

    search_dir = str(extracted_dirs[0]) if extracted_dirs else str(output_dir)

    # Security analysis
    print("[*] Searching for hardcoded credentials...")
    report["credentials"] = find_credentials(search_dir)

    print("[*] Identifying network services...")
    report["network_services"] = find_network_services(search_dir)

    print("[*] Finding executable binaries...")
    report["binaries"] = find_suspicious_binaries(search_dir)

    # Summary
    report["summary"] = {
        "signatures_found": len(report["binwalk"].get("signatures", [])),
        "credentials_found": len(report["credentials"]),
        "services_found": len(report["network_services"]),
        "binaries_found": len(report["binaries"]),
    }

    if args.format == "json":
        output = json.dumps(report, indent=2)
    else:
        output = f"=== Firmware Analysis Report ===\n"
        output += f"File: {report['firmware_file']}\n"
        output += f"SHA256: {report['firmware_hash']}\n"
        output += f"Size: {report['firmware_size']} bytes\n\n"
        s = report["summary"]
        output += f"Signatures: {s['signatures_found']} | Credentials: {s['credentials_found']} | "
        output += f"Services: {s['services_found']} | Binaries: {s['binaries_found']}\n\n"

        if report["credentials"]:
            output += "--- Credentials Found ---\n"
            for c in report["credentials"]:
                output += f"  [{c['severity'].upper()}] {c['type']} in {c['file']}\n"
            output += "\n"

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


if __name__ == "__main__":
    main()