# src/ctf/forensics.py
# Library for forensic analysis

import click
from dataclasses import dataclass, field
from pathlib import Path
import stat
import os
import sys
from typing import List, Dict

try:
    import pwd
    import grp
except ImportError:
    pwd = None
    grp = None

COMMON_SIGNATURES = {
    b"\x89PNG\r\n\x1a\n": ("PNG Image", [".png"]),
    b"\xff\xd8\xff": ("JPEG Image", [".jpg", ".jpeg"]),
    b"%PDF": ("PDF Document", [".pdf"]),
    b"PK\x03\x04": ("ZIP Archive", [".zip"]),
    b"\x7fELF": ("ELF Executable", [".elf"]),
    b"MZ": ("PE Executable", [".exe", ".dll"]),
    b"GIF87a": ("GIF Image", [".gif"]),
    b"GIF89a": ("GIF Image", [".gif"]),
    b"7z\xbc\xaf\x27\x1c": ("7-Zip Archive", [".7z"]),
    b"\x1f\x8b": ("GZIP Archive", [".gz"]),
    b"Rar!\x1a\x07\x00": ("RAR Archive", [".rar"]),
    b"Rar!\x1a\x07\x01\x00": ("RAR Archive", [".rar"]),
    b"BZh": ("BZIP2 Archive", [".bz2"]),
    b"BM": ("BMP Image", [".bmp"]),
    b"ID3": ("MP3 Audio", [".mp3"]),
}

@click.group(name="forensics")
def forensics_group(): pass

@forensics_group.command()
def tf(): print("hello from forensics")

@dataclass
class FileMetadata:
    filename: str
    size: int
    magic: str
    extension: str
    
    # Task 1: POSIX Permissions
    permissions_octal: str
    permissions_symbolic: str
    
    # Task 2: Ownership Identity
    owner_uid: int
    owner_username: str
    owner_gid: int
    owner_groupname: str
    
    # Task 3: Allocation Metrics
    allocated_size: int
    
    # Task 4: Hard Links
    hard_links: int
    
    # Task 5: Inode & Device Identifiers
    inode: int
    device: int
    
    # Task 6: Extended Attributes
    extended_attributes: Dict[str, str] = field(default_factory=dict)

@forensics_group.command()
@click.argument('path', type=click.Path(exists=True))
def inspect(path):
    p = Path(path)
    stat_info = p.stat()
    
    # Apparent size & extension
    size = stat_info.st_size
    extension = p.suffix
    
    # Magic bytes
    try:
        with open(p, 'rb') as f:
            magic = f.read(8).hex().upper()
    except Exception:
        magic = ""
        
    # POSIX Permissions
    mode = stat_info.st_mode
    permissions_octal = oct(stat.S_IMODE(mode))
    permissions_symbolic = stat.filemode(mode)
    
    # Ownership Identity
    owner_uid = stat_info.st_uid
    owner_gid = stat_info.st_gid
    owner_username = str(owner_uid)
    owner_groupname = str(owner_gid)
    
    if pwd is not None:
        try:
            owner_username = pwd.getpwuid(owner_uid).pw_name
        except KeyError:
            pass
    if grp is not None:
        try:
            owner_groupname = grp.getgrgid(owner_gid).gr_name
        except KeyError:
            pass
            
    # Allocation Metrics
    if hasattr(stat_info, "st_blocks"):
        allocated_size = stat_info.st_blocks * 512
    else:
        allocated_size = size
        
    # Hard Links
    hard_links = stat_info.st_nlink
    
    # Inode & Device Identifiers
    inode = stat_info.st_ino
    device = stat_info.st_dev
    
    # Extended Attributes
    extended_attributes = {}
    if hasattr(os, "listxattr") and hasattr(os, "getxattr"):
        try:
            attrs = os.listxattr(p)
            for attr in attrs:
                try:
                    val = os.getxattr(p, attr)
                    extended_attributes[attr] = val.decode("utf-8", errors="ignore")
                except OSError:
                    pass
        except OSError:
            pass
            
    meta = FileMetadata(
        filename=p.name,
        size=size,
        magic=magic,
        extension=extension,
        permissions_octal=permissions_octal,
        permissions_symbolic=permissions_symbolic,
        owner_uid=owner_uid,
        owner_username=owner_username,
        owner_gid=owner_gid,
        owner_groupname=owner_groupname,
        allocated_size=allocated_size,
        hard_links=hard_links,
        inode=inode,
        device=device,
        extended_attributes=extended_attributes
    )
    
    # Present using Rich Table
    from rich.console import Console
    from rich.table import Table
    
    console = Console()
    table = Table(title=f"Metadata: {meta.filename}", show_header=False)
    table.add_column("Key", style="bold cyan")
    table.add_column("Value")
    
    table.add_row("Filename", meta.filename)
    table.add_row("Size", f"{meta.size} bytes")
    table.add_row("Magic Bytes (Hex)", meta.magic)
    table.add_row("Extension", meta.extension)
    table.add_row("Permissions", f"{meta.permissions_symbolic} ({meta.permissions_octal})")
    table.add_row("Owner", f"{meta.owner_username} (UID: {meta.owner_uid})")
    table.add_row("Group", f"{meta.owner_groupname} (GID: {meta.owner_gid})")
    table.add_row("Allocated Space", f"{meta.allocated_size} bytes")
    table.add_row("Hard Links", str(meta.hard_links))
    table.add_row("Inode", str(meta.inode))
    table.add_row("Device", str(meta.device))
    
    if meta.extended_attributes:
        xattr_str = ", ".join(f"{k}={v}" for k, v in meta.extended_attributes.items())
        table.add_row("Extended Attributes", xattr_str)
        
    console.print(table)
    return meta

@forensics_group.command(name="signatures")
def list_signatures():
    """List all supported file magic signatures and expected extensions."""
    from rich.console import Console
    from rich.table import Table

    console = Console()
    table = Table(title="Supported Magic Signatures", show_header=True)
    table.add_column("Magic Bytes (Hex)", style="bold cyan")
    table.add_column("File Type", style="bold green")
    table.add_column("Expected Exts")

    for signature, (type_name, exts) in COMMON_SIGNATURES.items():
        hex_str = signature.hex().upper()
        exts_str = ", ".join(exts) if exts else "Any / None"
        table.add_row(hex_str, type_name, exts_str)

    console.print(table)