init
This commit is contained in:
senstella
198
pwr/apply.py
Normal file
198
pwr/apply.py
Normal file
@@ -0,0 +1,198 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user