diff --git a/.DS_Store b/.DS_Store index 2aab8c4..7a2552b 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 new file mode 100644 index 0000000..0b4cf31 Binary files /dev/null and b/__pycache__/base.cpython-314.pyc differ diff --git a/__pycache__/heic_types.cpython-314.pyc b/__pycache__/heic_types.cpython-314.pyc new file mode 100644 index 0000000..0aae52b Binary files /dev/null and b/__pycache__/heic_types.cpython-314.pyc differ diff --git a/__pycache__/parser.cpython-314.pyc b/__pycache__/parser.cpython-314.pyc index b2ca2b0..06faf08 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 new file mode 100644 index 0000000..f122d1a --- /dev/null +++ b/base.py @@ -0,0 +1,15 @@ +from typing import List + +class Box: + """ + Represents a generic ISOBMFF box. This is a simple data container. + """ + def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): + self.size = size + self.type = box_type + self.offset = offset + self.raw_data = raw_data + self.children: List['Box'] = [] # Initially empty + + def __repr__(self) -> str: + return f"" \ No newline at end of file diff --git a/heic_types.py b/heic_types.py new file mode 100644 index 0000000..739b7f1 --- /dev/null +++ b/heic_types.py @@ -0,0 +1,48 @@ +# In heic_types.py +import struct +from base import Box # <-- THIS IS THE KEY CHANGE + +class ItemLocation: + # ... (rest of this class is unchanged) + def __init__(self, item_id, offset, length): + self.item_id = item_id + self.offset = offset + self.length = length + def __repr__(self): + 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] + offset_size = (sizes >> 12) & 0x0F + length_size = (sizes >> 8) & 0x0F + base_offset_size = (sizes >> 4) & 0x0F + item_count = struct.unpack('>H', stream[6:8])[0] + current_pos = 8 + for _ in range(item_count): + item_id = struct.unpack('>H', stream[current_pos : current_pos+2])[0] + current_pos += 4 # Skip item_id and construction_method + current_pos += 2 # Skip data_reference_index + current_pos += base_offset_size + extent_count = struct.unpack('>H', stream[current_pos : current_pos+2])[0] + current_pos += 2 + if extent_count > 0: + offset = self._read_int(stream, current_pos, offset_size) + current_pos += offset_size + length = self._read_int(stream, current_pos, length_size) + current_pos += length_size + self.locations.append(ItemLocation(item_id, offset, length)) + def _read_int(self, data, pos, size): + if size == 1: return data[pos] + if size == 2: return struct.unpack('>H', data[pos:pos+2])[0] + if size == 4: return struct.unpack('>I', data[pos:pos+4])[0] + if size == 8: return struct.unpack('>Q', data[pos:pos+8])[0] + return 0 \ No newline at end of file diff --git a/main.py b/main.py index 5d017c6..b5dcc9f 100644 --- a/main.py +++ b/main.py @@ -1,13 +1,30 @@ +from base import Box from parser import parse_boxes +from heic_types import ItemLocationBox +from typing import List +import os + +def print_box_tree(boxes: List[Box], indent: int = 0): + for box in boxes: + print(" " * indent + str(box)) + + if isinstance(box, ItemLocationBox): + for loc in box.locations[:5]: + print(" " * (indent + 1) + str(loc)) + if len(box.locations) > 5: + print(" " * (indent + 1) + f"... and {len(box.locations) - 5} more locations") + + if box.children: + print_box_tree(box.children, indent + 1) def main(): # Test with the Apple HEIC file print("--- Analyzing Apple HEIC file ---") try: with open('apple.heic', 'rb') as f: - top_level_boxes = parse_boxes(f) - for box in top_level_boxes: - print(box) + file_size = os.fstat(f.fileno()).st_size + top_level_boxes = parse_boxes(f, file_size) + print_box_tree(top_level_boxes) except FileNotFoundError: print("Error: apple.heic not found.") @@ -17,9 +34,9 @@ def main(): print("--- Analyzing Samsung HEIC file ---") try: with open('samsung.heic', 'rb') as f: - top_level_boxes = parse_boxes(f) - for box in top_level_boxes: - print(box) + file_size = os.fstat(f.fileno()).st_size + top_level_boxes = parse_boxes(f, file_size) + print_box_tree(top_level_boxes) except FileNotFoundError: print("Error: samsung.heic not found.") diff --git a/parser.py b/parser.py index f701a23..26a0c76 100644 --- a/parser.py +++ b/parser.py @@ -1,78 +1,67 @@ import struct - -class Box: - """ - Represents a generic ISOBMFF box. - """ - def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes): - self.size = size - self.type = box_type - self.offset = offset - self.raw_data = raw_data - self.children = [] # Will be populated for container boxes - - def __repr__(self) -> str: - # e.g., - return f"" - -# In parser.py - from typing import List, BinaryIO +from io import BytesIO -def parse_boxes(stream: BinaryIO) -> List[Box]: +from base import Box # Import the base Box class +from heic_types import ItemLocationBox + +# The factory map remains here +BOX_TYPE_MAP = { + 'iloc': ItemLocationBox, +} + +CONTAINER_BOXES = {'meta', 'moov', 'trak', 'iprp', 'ipco', 'dinf', 'fiinf', 'ipro'} +FULL_BOXES = {'meta', 'hdlr', 'pitm', 'iinf', 'iloc'} + +def parse_boxes(stream: BinaryIO, max_size: int) -> List[Box]: """ - Parses all top-level boxes from a file stream, handling special size cases. + Parses boxes from a file stream up to a maximum size. + Now also handles recursive parsing. """ boxes = [] + start_pos_in_stream = stream.tell() - stream.seek(0, 2) - file_size = stream.tell() - stream.seek(0) - - while stream.tell() < file_size: - current_offset = stream.tell() + while stream.tell() - start_pos_in_stream < max_size: + current_offset_in_stream = stream.tell() - # Read the standard 8-byte header header = stream.read(8) - if len(header) < 8: - break + if len(header) < 8: break size = struct.unpack('>I', header[:4])[0] - box_type = header[4:].decode('ascii') + box_type = header[4:].decode('ascii', errors='ignore') header_size = 8 - content_size = 0 - - # --- NEW: Handle special size cases --- if size == 1: - # The actual size is a 64-bit integer following the type largesize_header = stream.read(8) - if len(largesize_header) < 8: - break + if len(largesize_header) < 8: break size = struct.unpack('>Q', largesize_header)[0] header_size = 16 - content_size = size - header_size elif size == 0: - # The box extends to the end of the file - content_size = file_size - current_offset - header_size - size = content_size + header_size - else: - # Standard size - content_size = size - header_size - # --- End of new code --- + size = max_size - (current_offset_in_stream - start_pos_in_stream) - # Read the box's content (raw_data) - # We need to rewind a bit if we read largesize header to get all content - stream.seek(current_offset + header_size) - raw_data = stream.read(content_size) - if len(raw_data) < content_size: - # Avoids errors on truncated files - break - - box = Box(size, box_type, current_offset, raw_data) - boxes.append(box) + if size < header_size: break + + content_size = size - header_size - # Seek to the beginning of the next box - stream.seek(current_offset + size) + stream.seek(current_offset_in_stream + header_size) + 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) + if box.type in FULL_BOXES: + 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) return boxes \ No newline at end of file