199 lines
6.8 KiB
Python
199 lines
6.8 KiB
Python
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)
|