[GEMINI] Write tests for forensics and stage baseline implementation
This commit is contained in:
@@ -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 ""
|
||||
# }}}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user