from __future__ import annotations import os from typing import BinaryIO from .formats import FilePatch, PatchReader from .proto import Control, SyncOp, SyncOpType from .proto import TlcContainer, TlcDir, TlcFile, TlcSymlink from .wire import BLOCK_SIZE class PatchApplyError(Exception): pass _MODE_MASK = 0o7777 class FilePool: def __init__(self, paths: list[str]): self._paths = list(paths) self._handles: dict[int, BinaryIO] = {} 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: handle = open(self._paths[index], "rb") self._handles[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: raise PatchApplyError("unexpected EOF while copying block range") dst.write(chunk) remaining -= len(chunk) def apply_rsync_ops(ops: list[SyncOp], target_pool: FilePool, output: BinaryIO) -> None: for op in ops: 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}") if op.file_index is None or op.block_index is None or op.block_span is None: raise PatchApplyError("missing fields in block range op") if op.block_span <= 0: raise PatchApplyError("invalid block span in block range op") file_size = target_pool.size(op.file_index) last_block_index = op.block_index + op.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 = (op.block_span - 1) * BLOCK_SIZE + last_block_size src = target_pool.open(op.file_index) src.seek(op.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: for entry in patch_reader.iter_file_entries(): if entry.sync_header.file_index is None: raise PatchApplyError("missing file_index in sync header") out_index = entry.sync_header.file_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) with open(out_path, "wb") as out: if entry.is_rsync(): if entry.sync_ops is None: raise PatchApplyError("missing rsync ops") apply_rsync_ops(entry.sync_ops, pool, out) elif entry.is_bsdiff(): if entry.bsdiff_header is None or entry.bsdiff_controls is None: raise PatchApplyError("missing bsdiff data") if entry.bsdiff_header.target_index is None: raise PatchApplyError("missing target_index in bsdiff header") old = pool.open(entry.bsdiff_header.target_index) apply_bsdiff_controls(entry.bsdiff_controls, old, out) else: 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)