[GEMINI] Write tests for forensics and stage baseline implementation

This commit is contained in:
venus
2026-07-18 03:28:52 -05:00
parent be58885e4f
commit 7bb0f87808
11 changed files with 574 additions and 34 deletions

View File

@@ -7,7 +7,9 @@ from pathlib import Path
import stat
import os
import sys
from typing import List, Dict
import struct
import xml.etree.ElementTree as ET
from typing import List, Dict, Any
try:
import pwd
@@ -16,6 +18,7 @@ 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"]),
@@ -33,6 +36,7 @@ COMMON_SIGNATURES = {
b"BM": ("BMP Image", [".bmp"]),
b"ID3": ("MP3 Audio", [".mp3"]),
}
# }}}
# {{{ FileMetadata
@dataclass
@@ -65,6 +69,12 @@ class FileMetadata:
# 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
@@ -146,6 +156,9 @@ def get_metadata(path: Path) -> FileMetadata:
except OSError:
pass
exif_data = get_exif(p)
comment = get_comment(p)
return FileMetadata(
filename=p.name,
size=size,
@@ -162,7 +175,344 @@ def get_metadata(path: Path) -> FileMetadata:
hard_links=hard_links,
inode=inode,
device=device,
extended_attributes=extended_attributes
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 ""
# }}}