# src/ctf/forensics.py
# Library for forensic analysis (pure functions only)

# vim foldmethod=marker
from dataclasses import dataclass, field
from pathlib import Path
import stat
import os
import sys
import struct
import xml.etree.ElementTree as ET
from typing import List, Dict, Any

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

# {{{ Common Signatures
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"]),
}
# }}}

# {{{ FileMetadata
@dataclass
class FileMetadata:
    filename: str
    size: int
    magic: str
    extension: str
    detected_type: 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)
    
    # EXIF Data
    exif_data: Dict[str, Any] = field(default_factory=dict)
    
    # JPEG/PNG Comments
    comment: str = ""
# }}}

# {{{ get_metadata
def get_metadata(path: Path) -> FileMetadata:
    """Extracts metadata attributes from a file without any console rendering."""
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"File not found: {p}")
        
    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 = ""
        
    detected_type = "Unknown"
    try:
        magic_bytes = bytes.fromhex(magic)
        for signature, (type_name, exts) in COMMON_SIGNATURES.items():
            if magic_bytes.startswith(signature):
                detected_type = type_name
                break
    except Exception:
        pass
        
    # 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
            
    exif_data = get_exif(p)
    comment = get_comment(p)
            
    return FileMetadata(
        filename=p.name,
        size=size,
        magic=magic,
        extension=extension,
        detected_type=detected_type,
        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,
        exif_data=exif_data,
        comment=comment
    )
# }}}

EXIF_TAGS = {
    0x010e: "ImageDescription",
    0x010f: "Make",
    0x0110: "Model",
    0x0112: "Orientation",
    0x011a: "XResolution",
    0x011b: "YResolution",
    0x0128: "ResolutionUnit",
    0x0131: "Software",
    0x0132: "DateTime",
    0x013b: "Artist",
    0x8298: "Copyright",
    0x829a: "ExposureTime",
    0x829d: "FNumber",
    0x8822: "ExposureProgram",
    0x8825: "GPSInfo",
    0x8827: "ISOSpeedRatings",
    0x9000: "ExifVersion",
    0x9003: "DateTimeOriginal",
    0x9004: "DateTimeDigitized",
    0x9201: "ShutterSpeedValue",
    0x9202: "ApertureValue",
    0x9204: "ExposureBiasValue",
    0x9205: "MaxApertureValue",
    0x9207: "MeteringMode",
    0x9208: "LightSource",
    0x9209: "Flash",
    0x920a: "FocalLength",
    0x9286: "UserComment",
    0xa000: "FlashpixVersion",
    0xa001: "ColorSpace",
    0xa002: "ExifImageWidth",
    0xa003: "ExifImageHeight",
    0xa402: "ExposureMode",
    0xa403: "WhiteBalance",
    0xa405: "FocalLengthIn35mmFilm",
    0xa406: "SceneCaptureType",
}

# {{{ parse_tiff
def parse_tiff(data: bytes) -> Dict[str, Any]:
    """Parses raw TIFF data block into EXIF tag dictionary."""
    tags = {}
    if len(data) < 8:
        return tags
        
    byte_order = data[0:2]
    if byte_order == b"II":
        endian = "<"
    elif byte_order == b"MM":
        endian = ">"
    else:
        return tags
        
    magic = struct.unpack(f"{endian}H", data[2:4])[0]
    if magic != 42:
        return tags
        
    ifd_offset = struct.unpack(f"{endian}I", data[4:8])[0]
    
    try:
        while ifd_offset != 0 and ifd_offset < len(data):
            if ifd_offset + 2 > len(data):
                break
            num_entries = struct.unpack(f"{endian}H", data[ifd_offset:ifd_offset+2])[0]
            entry_offset = ifd_offset + 2
            
            for _ in range(num_entries):
                if entry_offset + 12 > len(data):
                    break
                tag, field_type, count, val_offset = struct.unpack(
                    f"{endian}HHII", data[entry_offset:entry_offset+12]
                )
                entry_offset += 12
                
                # Resolve value based on type and count
                val = None
                type_sizes = {1: 1, 2: 1, 3: 2, 4: 4, 5: 8, 7: 1, 9: 4, 10: 8}
                size = type_sizes.get(field_type, 1) * count
                
                if size <= 4:
                    raw_val = struct.pack(f"{endian}I", val_offset)[:size]
                else:
                    if val_offset + size <= len(data):
                        raw_val = data[val_offset:val_offset+size]
                    else:
                        raw_val = b""
                        
                if field_type == 2:  # ASCII
                    val = raw_val.split(b"\x00")[0].decode("utf-8", errors="ignore")
                elif field_type == 3:  # SHORT
                    fmt = f"{endian}" + "H" * count
                    if len(raw_val) >= 2 * count:
                        val = struct.unpack(fmt, raw_val[:2*count])
                        if count == 1:
                            val = val[0]
                elif field_type == 4:  # LONG
                    fmt = f"{endian}" + "I" * count
                    if len(raw_val) >= 4 * count:
                        val = struct.unpack(fmt, raw_val[:4*count])
                        if count == 1:
                            val = val[0]
                elif field_type == 5:  # RATIONAL
                    fmt = f"{endian}" + "II" * count
                    if len(raw_val) >= 8 * count:
                        unpacked = struct.unpack(fmt, raw_val[:8*count])
                        rationals = [f"{unpacked[i*2]}/{unpacked[i*2+1]}" for i in range(count)]
                        val = rationals[0] if count == 1 else rationals
                elif field_type == 7:  # UNDEFINED
                    val = raw_val
                elif field_type == 9:  # SLONG
                    fmt = f"{endian}" + "i" * count
                    if len(raw_val) >= 4 * count:
                        val = struct.unpack(fmt, raw_val[:4*count])
                        if count == 1:
                            val = val[0]
                elif field_type == 10:  # SRATIONAL
                    fmt = f"{endian}" + "ii" * count
                    if len(raw_val) >= 8 * count:
                        unpacked = struct.unpack(fmt, raw_val[:8*count])
                        rationals = [f"{unpacked[i*2]}/{unpacked[i*2+1]}" for i in range(count)]
                        val = rationals[0] if count == 1 else rationals
                else:
                    val = raw_val
                    
                tag_name = EXIF_TAGS.get(tag, f"Tag_{hex(tag)}")
                tags[tag_name] = val
                
            if entry_offset + 4 > len(data):
                break
            ifd_offset = struct.unpack(f"{endian}I", data[entry_offset:entry_offset+4])[0]
    except Exception:
        pass
        
    return tags
# }}}

# {{{ extract_jpeg_exif
def extract_jpeg_exif(data: bytes) -> bytes:
    """Extracts raw Exif/TIFF payload from JPEG APP1 segments."""
    if not data.startswith(b"\xff\xd8\xff"):
        return b""
    idx = 2
    while idx < len(data) - 4:
        if data[idx] == 0xff:
            marker = data[idx+1]
            if marker == 0xd9:  # EOI
                break
            # Markers without length parameters
            if marker in (0xd8, 0xd9, 0x00) or 0xd0 <= marker <= 0xd7:
                idx += 2
                continue
            length = struct.unpack(">H", data[idx+2:idx+4])[0]
            if marker == 0xe1:  # APP1
                app1_data = data[idx+4:idx+4+length-2]
                if app1_data.startswith(b"Exif\x00\x00"):
                    return app1_data[6:]
            idx += 2 + length
        else:
            idx += 1
    return b""
# }}}

# {{{ extract_png_exif
def extract_png_exif(data: bytes) -> bytes:
    """Extracts raw Exif/TIFF payload from PNG eXIf chunks."""
    if not data.startswith(b"\x89PNG\r\n\x1a\n"):
        return b""
    idx = 8
    while idx < len(data) - 8:
        length = struct.unpack(">I", data[idx:idx+4])[0]
        chunk_type = data[idx+4:idx+8]
        if chunk_type == b"eXIf":
            return data[idx+8:idx+8+length]
        elif chunk_type == b"IEND":
            break
        idx += 12 + length
    return b""
# }}}

# {{{ get_exif
def get_exif(path: Path) -> Dict[str, Any]:
    """Reads file, extracts Exif and XMP segments, and parses them to a tag dictionary."""
    try:
        with open(path, "rb") as f:
            data = f.read()
    except Exception:
        return {}
        
    tags = {}
    
    # 1. Parse standard EXIF
    exif_data = b""
    if data.startswith(b"\xff\xd8\xff"):
        exif_data = extract_jpeg_exif(data)
    elif data.startswith(b"\x89PNG\r\n\x1a\n"):
        exif_data = extract_png_exif(data)
        
    if exif_data:
        try:
            tags.update(parse_tiff(exif_data))
        except Exception:
            pass
            
    # 2. Parse XMP (Adobe)
    xmp_str = ""
    if data.startswith(b"\xff\xd8\xff"):
        xmp_str = extract_jpeg_xmp(data)
        
    if xmp_str:
        try:
            tags.update(parse_xmp(xmp_str))
        except Exception:
            pass
            
    return tags
# }}}

# {{{ extract_jpeg_comment
def extract_jpeg_comment(data: bytes) -> str:
    """Extracts raw comment string from JPEG COM (0xfe) segments."""
    if not data.startswith(b"\xff\xd8\xff"):
        return ""
    idx = 2
    while idx < len(data) - 4:
        if data[idx] == 0xff:
            marker = data[idx+1]
            if marker == 0xd9:  # EOI
                break
            if marker in (0xd8, 0xd9, 0x00) or 0xd0 <= marker <= 0xd7:
                idx += 2
                continue
            length = struct.unpack(">H", data[idx+2:idx+4])[0]
            if marker == 0xfe:  # COM
                comment_data = data[idx+4:idx+4+length-2]
                return comment_data.decode("utf-8", errors="ignore")
            idx += 2 + length
        else:
            idx += 1
    return ""
# }}}

# {{{ get_comment
def get_comment(path: Path) -> str:
    """Reads file, checks headers, and extracts JPEG COM comments."""
    try:
        with open(path, "rb") as f:
            data = f.read()
    except Exception:
        return ""
        
    if data.startswith(b"\xff\xd8\xff"):
        return extract_jpeg_comment(data)
    return ""
# }}}

NS_MAP = {
    "http://creativecommons.org/ns#": "cc",
    "http://purl.org/dc/elements/1.1/": "dc",
    "http://ns.adobe.com/xap/1.0/": "xmp",
    "http://ns.adobe.com/exif/1.0/": "exif",
    "http://ns.adobe.com/tiff/1.0/": "tiff",
}

# {{{ parse_xmp
def parse_xmp(xmp_str: str) -> Dict[str, Any]:
    """Parses Adobe XMP XML string and returns namespace-prefixed tags."""
    metadata = {}
    try:
        start_idx = xmp_str.find("<x:xmpmeta")
        if start_idx == -1:
            return metadata
        end_idx = xmp_str.find("</x:xmpmeta>")
        if end_idx == -1:
            return metadata
        xml_data = xmp_str[start_idx:end_idx + len("</x:xmpmeta>")]
        
        root = ET.fromstring(xml_data)
        
        def get_clean_name(tag: str) -> str:
            if "}" in tag:
                ns, name = tag.split("}")
                ns = ns.lstrip("{")
                prefix = NS_MAP.get(ns, ns.split("/")[-1].split("#")[0])
                return f"{prefix}:{name}"
            return tag
            
        for elem in root.iter():
            clean_tag = get_clean_name(elem.tag)
            if clean_tag in ("rdf:RDF", "rdf:Description", "x:xmpmeta"):
                continue
                
            val = None
            if elem.text and elem.text.strip():
                val = elem.text.strip()
            else:
                attribs = {get_clean_name(k): v for k, v in elem.attrib.items()}
                attribs.pop("rdf:about", None)
                if "rdf:resource" in attribs:
                    val = attribs["rdf:resource"]
                elif attribs:
                    val = ", ".join(f"{k}={v}" for k, v in attribs.items())
                    
            if val:
                metadata[clean_tag] = val
    except Exception:
        pass
    return metadata
# }}}

# {{{ extract_jpeg_xmp
def extract_jpeg_xmp(data: bytes) -> str:
    """Extracts XMP XML string from JPEG APP1 segments."""
    if not data.startswith(b"\xff\xd8\xff"):
        return ""
    idx = 2
    while idx < len(data) - 4:
        if data[idx] == 0xff:
            marker = data[idx+1]
            if marker == 0xd9:  # EOI
                break
            if marker in (0xd8, 0xd9, 0x00) or 0xd0 <= marker <= 0xd7:
                idx += 2
                continue
            length = struct.unpack(">H", data[idx+2:idx+4])[0]
            if marker == 0xe1:  # APP1
                app1_data = data[idx+4:idx+4+length-2]
                if app1_data.startswith(b"http://ns.adobe.com/xap/1.0/\x00"):
                    return app1_data[29:].decode("utf-8", errors="ignore")
            idx += 2 + length
        else:
            idx += 1
    return ""
# }}}

