from __future__ import annotations import errno import os from collections import OrderedDict from typing import BinaryIO from .formats import PatchReader from .proto import Control, SyncOp, SyncOpType, TlcContainer, TlcFile from .wire import BLOCK_SIZE class PatchApplyError(Exception): pass _MODE_MASK = 0o7777 class FilePool: def __init__(self, paths: list[str], max_open: int = 128): self._paths = list(paths) self._handles: OrderedDict[int, BinaryIO] = OrderedDict() self._max_open = max(1, int(max_open)) def _touch(self, index: int, handle: BinaryIO) -> None: self._handles.pop(index, None) self._handles[index] = handle def _evict_if_needed(self) -> None: while len(self._handles) >= self._max_open: _, handle = self._handles.popitem(last=False) handle.close() def open(self, index: int) -> BinaryIO: if index < 0 or index >= len(self._paths): raise PatchApplyError(f"file index out of range: {index}") handle = self._handles.get(index) if handle is None: try: self._evict_if_needed() handle = open(self._paths[index], "rb") except OSError as exc: if exc.errno in (errno.EMFILE, errno.ENFILE): self.close() handle = open(self._paths[index], "rb") else: raise self._handles[index] = handle else: self._touch(index, handle) return handle def size(self, index: int) -> int: if index < 0 or index >= len(self._paths): raise PatchApplyError(f"file index out of range: {index}") return os.path.getsize(self._paths[index]) def close(self) -> None: for handle in self._handles.values(): handle.close() self._handles.clear() def _copy_range( dst: BinaryIO, src: BinaryIO, length: int, buffer_size: int = 32 * 1024 ) -> None: remaining = length while remaining > 0: chunk = src.read(min(buffer_size, remaining)) if not chunk: break dst.write(chunk) remaining -= len(chunk) def _copy_all(dst: BinaryIO, src: BinaryIO, buffer_size: int = 32 * 1024) -> None: while True: chunk = src.read(buffer_size) if not chunk: return dst.write(chunk) def _compute_num_blocks(file_size: int) -> int: return (file_size + BLOCK_SIZE - 1) // BLOCK_SIZE def _normalize_op_type(op_type: SyncOpType | int | None) -> SyncOpType | int: return op_type if op_type is not None else SyncOpType.BLOCK_RANGE def _is_full_file_op(op: SyncOp, target_size: int, output_size: int) -> bool: op_type = _normalize_op_type(op.type) if op_type != SyncOpType.BLOCK_RANGE: return False block_index = 0 if op.block_index is None else op.block_index if block_index != 0: return False if target_size != output_size: return False block_span = 0 if op.block_span is None else op.block_span return block_span == _compute_num_blocks(output_size) def _apply_file_mode(path: str, file: TlcFile | None) -> None: if file is None or file.mode is None: return mode = int(file.mode) & _MODE_MASK try: os.chmod(path, mode) except (PermissionError, OSError): pass def apply_rsync_ops(ops: list[SyncOp], target_pool: FilePool, output: BinaryIO) -> None: for op in ops: op_type = _normalize_op_type(op.type) if op_type == SyncOpType.DATA: output.write(op.data or b"") continue if op_type != SyncOpType.BLOCK_RANGE: raise PatchApplyError(f"unsupported sync op type: {op_type}") file_index = 0 if op.file_index is None else op.file_index block_index = 0 if op.block_index is None else op.block_index block_span = 0 if op.block_span is None else op.block_span file_size = target_pool.size(file_index) last_block_index = block_index + block_span - 1 last_block_size = BLOCK_SIZE if BLOCK_SIZE * (last_block_index + 1) > file_size: last_block_size = file_size % BLOCK_SIZE op_size = (block_span - 1) * BLOCK_SIZE + last_block_size src = target_pool.open(file_index) src.seek(block_index * BLOCK_SIZE) _copy_range(output, src, op_size) def apply_bsdiff_controls( controls: list[Control], old: BinaryIO, output: BinaryIO ) -> None: old_offset = 0 for ctrl in controls: if ctrl.eof: break add = ctrl.add or b"" copy = ctrl.copy or b"" seek = ctrl.seek or 0 if add: old.seek(old_offset) old_chunk = old.read(len(add)) if len(old_chunk) != len(add): raise PatchApplyError("unexpected EOF while reading bsdiff add data") out = bytes(((a + b) & 0xFF) for a, b in zip(old_chunk, add)) output.write(out) old_offset += len(add) if copy: output.write(copy) old_offset += seek def _ensure_parent(path: str) -> None: parent = os.path.dirname(path) if parent: os.makedirs(parent, exist_ok=True) def apply_patch( patch_reader: PatchReader, target_paths: list[str], output_paths: list[str] ) -> None: pool = FilePool(target_paths) try: expected_files = len(output_paths) if patch_reader.source_container is not None: expected_files = len(patch_reader.source_container.files) entry_iter = patch_reader.iter_file_entries() for expected_index in range(expected_files): try: entry = next(entry_iter) except StopIteration: raise PatchApplyError( f"corrupted patch: expected {expected_files} file entries, got {expected_index}" ) header_index = ( 0 if entry.sync_header.file_index is None else int(entry.sync_header.file_index) ) if header_index != expected_index: raise PatchApplyError( f"corrupted patch: expected file index {expected_index}, got {header_index}" ) out_index = header_index if out_index < 0 or out_index >= len(output_paths): raise PatchApplyError(f"output index out of range: {out_index}") out_path = output_paths[out_index] _ensure_parent(out_path) if entry.is_rsync(): if entry.sync_ops is None: raise PatchApplyError("missing rsync ops") if entry.sync_ops: op = entry.sync_ops[0] target_index = 0 if op.file_index is None else op.file_index target_file = None output_file = None if patch_reader.target_container and patch_reader.source_container: if 0 <= target_index < len(patch_reader.target_container.files): target_file = patch_reader.target_container.files[ target_index ] if 0 <= out_index < len(patch_reader.source_container.files): output_file = patch_reader.source_container.files[out_index] if target_file is not None and output_file is not None: target_size = ( int(target_file.size) if target_file.size is not None else 0 ) output_size = ( int(output_file.size) if output_file.size is not None else 0 ) if _is_full_file_op(op, target_size, output_size): src = pool.open(target_index) src.seek(0) with open(out_path, "wb") as out: _copy_all(out, src) _apply_file_mode(out_path, output_file) continue with open(out_path, "wb") as out: apply_rsync_ops(entry.sync_ops, pool, out) output_file = None if patch_reader.source_container and 0 <= out_index < len( patch_reader.source_container.files ): output_file = patch_reader.source_container.files[out_index] _apply_file_mode(out_path, output_file) continue if entry.is_bsdiff(): if entry.bsdiff_header is None or entry.bsdiff_controls is None: raise PatchApplyError("missing bsdiff data") target_index = ( 0 if entry.bsdiff_header.target_index is None else entry.bsdiff_header.target_index ) with open(out_path, "wb") as out: old = pool.open(target_index) apply_bsdiff_controls(entry.bsdiff_controls, old, out) expected_size = None output_file = None if patch_reader.source_container and 0 <= out_index < len( patch_reader.source_container.files ): output_file = patch_reader.source_container.files[out_index] if output_file.size is not None: expected_size = int(output_file.size) final_size = out.tell() if expected_size is not None and final_size != expected_size: raise PatchApplyError( f"corrupted patch: expected output size {expected_size}, got {final_size}" ) _apply_file_mode(out_path, output_file) continue raise PatchApplyError("unknown file patch type") finally: pool.close() def _container_file_paths(container: TlcContainer, root: str) -> list[str]: paths = [] for f in container.files: if f.path is None: raise PatchApplyError("container file missing path") paths.append(os.path.join(root, f.path)) return paths def _ensure_dirs(container: TlcContainer, root: str) -> None: for d in container.dirs: if d.path is None: continue path = os.path.join(root, d.path) os.makedirs(path, exist_ok=True) if d.mode is not None: mode = int(d.mode) & _MODE_MASK try: os.chmod(path, mode) except (PermissionError, OSError): pass def _ensure_symlinks(container: TlcContainer, root: str) -> None: for s in container.symlinks: if s.path is None or s.dest is None: continue path = os.path.join(root, s.path) if os.path.lexists(path): continue parent = os.path.dirname(path) if parent: os.makedirs(parent, exist_ok=True) try: os.symlink(s.dest, path) except (AttributeError, OSError): continue def apply_patch_to_folders(patch_path: str, target_root: str, output_root: str) -> None: with PatchReader.open( patch_path, container_decoder=(TlcContainer.from_bytes, TlcContainer.from_bytes), ) as reader: if reader.target_container is None or reader.source_container is None: raise PatchApplyError("missing containers in patch") target_container: TlcContainer = reader.target_container source_container: TlcContainer = reader.source_container _ensure_dirs(source_container, output_root) _ensure_symlinks(source_container, output_root) target_paths = _container_file_paths(target_container, target_root) output_paths = _container_file_paths(source_container, output_root) apply_patch(reader, target_paths, output_paths)