Nymiro
96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
import os
|
|
from base import Box
|
|
from parser import parse_boxes
|
|
from heic_types import ItemLocationBox
|
|
from heic_types import ItemLocationBox, PrimaryItemBox, ItemInfoBox
|
|
|
|
# Import handlers
|
|
from handlers.base_handler import VendorHandler
|
|
from handlers.apple_handler import AppleHandler
|
|
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
|
|
|
|
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()
|
|
|
|
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)
|
|
|
|
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.")
|
|
return
|
|
|
|
print("Available items in HEIC file:")
|
|
for entry in self._iinf_box.entries:
|
|
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]:
|
|
self.handler = SamsungHandler()
|
|
return
|
|
if 'apple' in compatible_brands.lower() or 'MiHB' in compatible_brands:
|
|
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
|
|
location = next((loc for loc in self._iloc_box.locations if loc.item_id == item_id), None)
|
|
if not location:
|
|
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.
|
|
"""
|
|
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 |