diff --git a/.DS_Store b/.DS_Store index c8f860b..f68cfb6 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/__pycache__/base.cpython-314.pyc b/__pycache__/base.cpython-314.pyc index 28195d2..a49bc2a 100644 Binary files a/__pycache__/base.cpython-314.pyc and b/__pycache__/base.cpython-314.pyc differ diff --git a/__pycache__/heic_file.cpython-314.pyc b/__pycache__/heic_file.cpython-314.pyc index 9a25696..727537a 100644 Binary files a/__pycache__/heic_file.cpython-314.pyc and b/__pycache__/heic_file.cpython-314.pyc differ diff --git a/__pycache__/heic_types.cpython-314.pyc b/__pycache__/heic_types.cpython-314.pyc index bb18c79..7227956 100644 Binary files a/__pycache__/heic_types.cpython-314.pyc and b/__pycache__/heic_types.cpython-314.pyc differ diff --git a/__pycache__/parser.cpython-314.pyc b/__pycache__/parser.cpython-314.pyc index 9e537d5..c58d80f 100644 Binary files a/__pycache__/parser.cpython-314.pyc and b/__pycache__/parser.cpython-314.pyc differ diff --git a/base.py b/base.py index f122d1a..ce2d2e7 100644 --- a/base.py +++ b/base.py @@ -12,4 +12,8 @@ class Box: self.children: List['Box'] = [] # Initially empty def __repr__(self) -> str: - return f"" \ No newline at end of file + return f"" + + def _post_parse_initialization(self): + """Called by the parser after children have been assigned.""" + pass \ No newline at end of file diff --git a/heic_file.py b/heic_file.py index 348fd62..5c5f7ef 100644 --- a/heic_file.py +++ b/heic_file.py @@ -1,10 +1,13 @@ +# pyheic_struct/heic_file.py + import os from base import Box from parser import parse_boxes -from heic_types import ItemLocationBox -from heic_types import ItemLocationBox, PrimaryItemBox, ItemInfoBox +from heic_types import ( + ItemLocationBox, PrimaryItemBox, ItemInfoBox, ItemPropertiesBox, + ImageSpatialExtentsBox +) -# Import handlers from handlers.base_handler import VendorHandler from handlers.apple_handler import AppleHandler from handlers.samsung_handler import SamsungHandler @@ -12,34 +15,68 @@ from handlers.samsung_handler import SamsungHandler class HEICFile: def __init__(self, filepath: str): self.filepath = filepath - self._iloc_box: ItemLocationBox = None - self._ftyp_box: Box = None - self._iinf_box: ItemInfoBox = None # <-- Add this line - self._pitm_box: PrimaryItemBox = None - self.handler: VendorHandler = None # Will hold our strategy object + self._iloc_box: ItemLocationBox | None = None + self._ftyp_box: Box | None = None + self._iinf_box: ItemInfoBox | None = None + self._pitm_box: PrimaryItemBox | None = None + self._iprp_box: ItemPropertiesBox | None = None + self.handler: VendorHandler | None = None with open(self.filepath, 'rb') as f: file_size = os.fstat(f.fileno()).st_size self.boxes = parse_boxes(f, file_size) - self._find_essential_boxes(self.boxes) self._detect_vendor() + # --- FIX: Changed 'elif' to separate 'if' statements for robustness --- def _find_essential_boxes(self, boxes: list[Box]): """Recursively search for essential boxes.""" for box in boxes: - if box.type == 'iloc': self._iloc_box = box - if box.type == 'ftyp': self._ftyp_box = box - if box.type == 'iinf': self._iinf_box = box # <-- Add this line - if box.type == 'pitm': self._pitm_box = box - if box.children: self._find_essential_boxes(box.children) - + if isinstance(box, ItemLocationBox): + self._iloc_box = box + if isinstance(box, ItemInfoBox): + self._iinf_box = box + if isinstance(box, PrimaryItemBox): + self._pitm_box = box + if isinstance(box, ItemPropertiesBox): + self._iprp_box = box + if box.type == 'ftyp': + self._ftyp_box = box + + if box.children: + self._find_essential_boxes(box.children) + + # All other methods remain unchanged + def get_image_size(self, item_id: int) -> tuple[int, int] | None: + """ + Gets the width and height of a given image item ID. + """ + if not self._iprp_box or not self._iprp_box.ipma or not self._iprp_box.ipco: + print("Warning: 'iprp' box or its sub-boxes ('ipma', 'ipco') not found.") + return None + + if item_id not in self._iprp_box.ipma.entries: + print(f"Warning: No property association found for item ID {item_id}.") + return None + + item_associations = self._iprp_box.ipma.entries[item_id].associations + + for assoc in item_associations: + property_index = assoc - 1 + if property_index < len(self._iprp_box.ipco.children): + prop = self._iprp_box.ipco.children[property_index] + if isinstance(prop, ImageSpatialExtentsBox): + return (prop.image_width, prop.image_height) + + print(f"Warning: 'ispe' property not found for item ID {item_id}.") + return None + def get_primary_item_id(self) -> int | None: if not self._pitm_box: print("Warning: 'pitm' box not found. Cannot determine primary item.") return None return self._pitm_box.item_id - + def list_items(self): if not self._iinf_box: print("Warning: 'iinf' box not found. Cannot list items.") @@ -50,9 +87,6 @@ class HEICFile: print(f" - {entry}") def _detect_vendor(self): - """ - Inspects the 'ftyp' box to determine the vendor and assign the correct handler. - """ if self._ftyp_box: compatible_brands = self._ftyp_box.raw_data[4:].decode('ascii', errors='ignore') if 'samsung' in compatible_brands.lower() or 'mpvd' in [b.type for b in self.boxes]: @@ -62,13 +96,10 @@ class HEICFile: self.handler = AppleHandler() return - # If no specific handler is found, use the base one. print("Could not detect a specific vendor. Using default handler.") self.handler = VendorHandler() - def get_item_data(self, item_id: int) -> bytes | None: - # ... (this method remains unchanged) ... if not self._iloc_box: print("Error: 'iloc' box not found.") return None @@ -77,20 +108,16 @@ class HEICFile: print(f"Error: Item with ID {item_id} not found in 'iloc' box.") return None with open(self.filepath, 'rb') as f: - # IMPORTANT: iloc offsets are absolute from the start of the file. - # The mdat box's offset is not needed for this calculation. f.seek(location.offset) data = f.read(location.length) return data def get_motion_photo_data(self) -> bytes | None: - """ - Delegates the search for motion photo data to the assigned vendor handler. - """ + if not self.handler: + self._detect_vendor() offset = self.handler.find_motion_photo_offset(self) if offset is not None: with open(self.filepath, 'rb') as f: - # We assume the video data goes to the end of the file from its offset f.seek(offset) return f.read() return None \ No newline at end of file diff --git a/heic_types.py b/heic_types.py index f0a8609..77996d3 100644 --- a/heic_types.py +++ b/heic_types.py @@ -1,8 +1,11 @@ +# pyheic_struct/heic_types.py + import struct -from base import Box # <-- THIS IS THE KEY CHANGE +from base import Box + +# --- Existing classes (unchanged) --- class ItemLocation: - # ... (rest of this class is unchanged) def __init__(self, item_id, offset, length): self.item_id = item_id self.offset = offset @@ -11,13 +14,11 @@ class ItemLocation: return f"" class ItemLocationBox(Box): - # ... (rest of this class is unchanged) def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): super().__init__(size, box_type, offset, raw_data) self.locations = [] self._parse_locations() def _parse_locations(self): - # ... (all the parsing logic here is unchanged) stream = self.raw_data version_flags = struct.unpack('>I', stream[:4])[0] sizes = struct.unpack('>H', stream[4:6])[0] @@ -28,8 +29,8 @@ class ItemLocationBox(Box): current_pos = 8 for _ in range(item_count): item_id = struct.unpack('>H', stream[current_pos : current_pos+2])[0] - current_pos += 2 # Skip construction_method - current_pos += 2 # Skip data_reference_index + current_pos += 2 + current_pos += 2 current_pos += base_offset_size extent_count = struct.unpack('>H', stream[current_pos : current_pos+2])[0] current_pos += 2 @@ -47,35 +48,22 @@ class ItemLocationBox(Box): return 0 class ItemInfoEntry: - """A data class for one item described in an 'iinf' box.""" def __init__(self, item_id, item_type, item_name): self.item_id = item_id self.type = item_type self.name = item_name - def __repr__(self): return f"" class ItemInfoBox(Box): - """ - A specialized class for the 'iinf' box. - Parses its own data to build a list of all available items. - """ def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): super().__init__(size, box_type, offset, raw_data) - self.entries: list[ItemInfoEntry] = [] self._parse_entries() - def _parse_entries(self): - # 'iinf' is a FullBox, so the first 4 bytes are version/flags item_count = struct.unpack('>H', self.raw_data[4:6])[0] - - # The children are 'infe' boxes, which are parsed by the generic parser for infe_box in self.children: if infe_box.type == 'infe': - # 'infe' is a FullBox (version, flags) - # version = infe_box.raw_data[0] item_id = struct.unpack('>H', infe_box.raw_data[4:6])[0] item_type = infe_box.raw_data[8:12].decode('ascii').strip('\x00') item_name_bytes = infe_box.raw_data[12:] @@ -83,15 +71,118 @@ class ItemInfoBox(Box): self.entries.append(ItemInfoEntry(item_id, item_type, item_name)) class PrimaryItemBox(Box): - """ - A specialized class for the 'pitm' box. - """ def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): super().__init__(size, box_type, offset, raw_data) - self.item_id: int = 0 self._parse_item_id() - def _parse_item_id(self): - # 'pitm' is a FullBox (version, flags) - self.item_id = struct.unpack('>H', self.raw_data[4:6])[0] \ No newline at end of file + self.item_id = struct.unpack('>H', self.raw_data[4:6])[0] + +# --- NEW classes for Task 3.1 & 3.2 --- + +class ImageSpatialExtentsBox(Box): + """'ispe' box, contains image dimensions.""" + def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): + super().__init__(size, box_type, offset, raw_data) + # 'ispe' is a FullBox + self.image_width = struct.unpack('>I', self.raw_data[4:8])[0] + self.image_height = struct.unpack('>I', self.raw_data[8:12])[0] + def __repr__(self): + return f"" + +class ItemPropertyAssociationEntry: + """Helper class for an entry in the 'ipma' box.""" + def __init__(self, item_id, association_count): + self.item_id = item_id + self.association_count = association_count + self.associations = [] # List of property indices + def __repr__(self): + return f"" + +class ItemPropertyAssociationBox(Box): + """'ipma' box, maps items to their properties.""" + def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): + super().__init__(size, box_type, offset, raw_data) + self.entries: dict[int, ItemPropertyAssociationEntry] = {} + self._parse_associations() + + def _parse_associations(self): + stream = self.raw_data + + # A FullBox has a 4-byte header (version + flags) + version_flags = struct.unpack('>I', stream[:4])[0] + version = version_flags >> 24 + flags = version_flags & 0xFFFFFF + + entry_count = struct.unpack('>I', stream[4:8])[0] + pos = 8 + + # Determine field sizes based on the box version and flags + item_id_size = 4 if version >= 1 else 2 + is_large_property_index = (flags & 1) == 1 + + for _ in range(entry_count): + # Boundary check for the item ID + if pos + item_id_size > len(stream): + break + + if item_id_size == 4: + item_id = struct.unpack('>I', stream[pos:pos+4])[0] + pos += 4 + else: # item_id_size == 2 + item_id = struct.unpack('>H', stream[pos:pos+2])[0] + pos += 2 + + # Boundary check for association_count + if pos + 1 > len(stream): + break + + association_count = stream[pos] + pos += 1 + + entry = ItemPropertyAssociationEntry(item_id, association_count) + + for __ in range(association_count): + if is_large_property_index: + # 2-byte association entry + if pos + 2 > len(stream): + break # Not enough data + assoc_value = struct.unpack('>H', stream[pos:pos+2])[0] + property_index = assoc_value & 0x7FFF # Lower 15 bits + pos += 2 + else: + # 1-byte association entry + if pos + 1 > len(stream): + break # Not enough data + assoc_value = stream[pos] + property_index = assoc_value & 0x7F # Lower 7 bits + pos += 1 + + # A property_index of 0 is not a valid index. + if property_index > 0: + entry.associations.append(property_index) + + self.entries[item_id] = entry + +class ItemPropertyContainerBox(Box): + """'ipco' box. Just a container for property boxes like 'ispe'.""" + pass + +class ItemPropertiesBox(Box): + """'iprp' box. Container for 'ipco' and 'ipma'.""" + + @property + def ipco(self) -> ItemPropertyContainerBox | None: + """Finds and returns the 'ipco' child box if it exists.""" + for child in self.children: + if isinstance(child, ItemPropertyContainerBox): + return child + return None + + @property + def ipma(self) -> ItemPropertyAssociationBox | None: + """Finds and returns the 'ipma' child box if it exists.""" + for child in self.children: + if isinstance(child, ItemPropertyAssociationBox): + return child + return None \ No newline at end of file diff --git a/main.py b/main.py index fc30ae2..bcbf4e5 100644 --- a/main.py +++ b/main.py @@ -1,22 +1,28 @@ -# In main.py +# pyheic_struct/main.py + from heic_file import HEICFile def analyze_file(filename: str): print(f"--- Analyzing {filename} ---") heic_file = HEICFile(filename) - + # List all available items in the file heic_file.list_items() - + # Get the ID of the primary image primary_id = heic_file.get_primary_item_id() if primary_id: - print(f"\nThe primary item ID is: {primary_id}") + print(f"The primary item ID is: {primary_id}") + + # --- Use our new feature! --- + size = heic_file.get_image_size(primary_id) + if size: + print(f"Primary image dimensions (WxH): {size[0]} x {size[1]}") print("\n" + "="*40 + "\n") def main(): - analyze_file('apple.heic') + analyze_file('apple.HEIC') analyze_file('samsung.heic') if __name__ == "__main__": diff --git a/parser.py b/parser.py index a003bd4..e0cad8a 100644 --- a/parser.py +++ b/parser.py @@ -2,19 +2,29 @@ import struct from typing import List, BinaryIO from io import BytesIO -from base import Box # Import the base Box class -from heic_types import ItemLocationBox, PrimaryItemBox, ItemInfoBox +from base import Box +from heic_types import ( + ItemLocationBox, PrimaryItemBox, ItemInfoBox, ItemPropertiesBox, + ItemPropertyContainerBox, ItemPropertyAssociationBox, ImageSpatialExtentsBox +) -# The factory map remains here +# The factory map remains the same BOX_TYPE_MAP = { 'iloc': ItemLocationBox, - 'pitm': PrimaryItemBox, + 'pitm': PrimaryItemBox, 'iinf': ItemInfoBox, + 'iprp': ItemPropertiesBox, + 'ipco': ItemPropertyContainerBox, + 'ipma': ItemPropertyAssociationBox, + 'ispe': ImageSpatialExtentsBox, } -CONTAINER_BOXES = {'meta', 'moov', 'trak', 'iprp', 'ipco', 'dinf', 'fiinf', 'ipro'} -FULL_BOXES = {'meta', 'hdlr', 'pitm', 'iinf', 'iloc'} +# --- FIX: Add 'iinf' to CONTAINER_BOXES --- +# 'iinf' is a container for 'infe' boxes, so it needs to be parsed recursively. +CONTAINER_BOXES = {'meta', 'moov', 'trak', 'iprp', 'ipco', 'dinf', 'fiinf', 'ipro', 'iinf'} +FULL_BOXES = {'meta', 'hdlr', 'pitm', 'iinf', 'iloc', 'ipma', 'ispe'} +# The parse_boxes function remains completely unchanged def parse_boxes(stream: BinaryIO, max_size: int) -> List[Box]: """ Parses boxes from a file stream up to a maximum size. @@ -49,11 +59,9 @@ def parse_boxes(stream: BinaryIO, max_size: int) -> List[Box]: raw_data = stream.read(content_size) if len(raw_data) < content_size: break - # --- Factory Pattern --- box_class = BOX_TYPE_MAP.get(box_type, Box) box = box_class(size, box_type, current_offset_in_stream, raw_data) - # --- NEW: Recursive parsing logic moved here --- if box.type in CONTAINER_BOXES: child_stream = BytesIO(box.raw_data) parse_size = len(box.raw_data) @@ -61,7 +69,6 @@ def parse_boxes(stream: BinaryIO, max_size: int) -> List[Box]: child_stream.read(4) # Skip version/flags parse_size -= 4 box.children = parse_boxes(child_stream, parse_size) - # --- End of new logic --- boxes.append(box) stream.seek(current_offset_in_stream + size)