108 lines
4.4 KiB
Python
108 lines
4.4 KiB
Python
import copy
|
|
import os
|
|
import zipfile
|
|
import hashlib
|
|
import tempfile
|
|
import json
|
|
from collections import OrderedDict
|
|
from threading import Lock
|
|
|
|
from .errors import PackageReadError
|
|
from .limits import PackageLimits
|
|
from .validator import ARCHIVE_ERRORS, PackageValidator
|
|
|
|
_VALIDATED = OrderedDict()
|
|
_CACHE_LOCK = Lock()
|
|
|
|
|
|
class PackageReader:
|
|
"""Read-only, validated access; no extraction or package rewriting."""
|
|
|
|
def __init__(self, path: str | os.PathLike[str], limits: PackageLimits | None = None):
|
|
self.path = path
|
|
self.validator = PackageValidator(limits)
|
|
self._source = None
|
|
self._archive = None
|
|
self._result = None
|
|
|
|
def __enter__(self):
|
|
if self._source is not None:
|
|
raise RuntimeError("Reader is already open")
|
|
try:
|
|
self._source = tempfile.SpooledTemporaryFile(max_size=8*1024*1024,mode="w+b")
|
|
total = 0
|
|
archive_hash = hashlib.sha256()
|
|
# The archive used for reads is the exact bounded snapshot whose
|
|
# bytes were hashed. External file changes cannot invalidate it.
|
|
with open(self.path,"rb") as incoming:
|
|
if os.fstat(incoming.fileno()).st_size > self.validator.limits.max_archive_bytes:
|
|
raise PackageReadError(self.validator._failure("ARCHIVE_LIMIT", "Archive exceeds configured size limit"))
|
|
while chunk := incoming.read(65536):
|
|
total += len(chunk)
|
|
if total > self.validator.limits.max_archive_bytes:
|
|
raise PackageReadError(self.validator._failure("ARCHIVE_LIMIT", "Archive exceeds configured size limit"))
|
|
archive_hash.update(chunk)
|
|
self._source.write(chunk)
|
|
self._source.seek(0)
|
|
self._archive = zipfile.ZipFile(self._source)
|
|
cache_key = (repr(self.validator.limits),archive_hash.digest())
|
|
with _CACHE_LOCK:
|
|
cached = _VALIDATED.get(cache_key)
|
|
if cached is not None:
|
|
_VALIDATED.move_to_end(cache_key)
|
|
self._result = copy.deepcopy(cached[0])
|
|
if self._result is None:
|
|
self._result = self.validator._validate_archive(self._archive)
|
|
if not self._result.valid:
|
|
raise PackageReadError(self._result)
|
|
if cached is None:
|
|
# Conservative metadata weight; oversized results are simply
|
|
# not cached. Bound both entry count and aggregate footprint.
|
|
weight=len(json.dumps(self._result.manifest).encode())*32+sum(len(name)*4+768 for name in self._result.files)
|
|
if weight<=16*1024*1024:
|
|
with _CACHE_LOCK:
|
|
_VALIDATED[cache_key] = (copy.deepcopy(self._result),weight)
|
|
while len(_VALIDATED)>32 or sum(value[1] for value in _VALIDATED.values())>16*1024*1024:
|
|
_VALIDATED.popitem(last=False)
|
|
return self
|
|
except PackageReadError:
|
|
self.close()
|
|
raise
|
|
except ARCHIVE_ERRORS as error:
|
|
self.close()
|
|
raise PackageReadError(self.validator._failure("ARCHIVE_ERROR", str(error))) from error
|
|
|
|
def _require_open(self):
|
|
if self._archive is None or self._result is None:
|
|
raise RuntimeError("Use PackageReader inside a with block")
|
|
|
|
@property
|
|
def manifest(self):
|
|
self._require_open()
|
|
return copy.deepcopy(self._result.manifest)
|
|
|
|
@property
|
|
def files(self) -> tuple[str, ...]:
|
|
self._require_open()
|
|
return self._result.files
|
|
|
|
def read_bytes(self, path: str) -> bytes:
|
|
self._require_open()
|
|
if path not in self._result.files:
|
|
raise PackageReadError(self.validator._failure("FILE_NOT_FOUND", "File is not in package", path))
|
|
try:
|
|
with self._archive.open(path) as stream:
|
|
return self.validator._read_bounded(stream, self.validator.limits.max_file_bytes)
|
|
except ARCHIVE_ERRORS as error:
|
|
raise PackageReadError(self.validator._failure("READ_ERROR", str(error), path)) from error
|
|
|
|
def close(self):
|
|
if self._archive is not None:
|
|
self._archive.close()
|
|
if self._source is not None:
|
|
self._source.close()
|
|
self._archive = self._source = self._result = None
|
|
|
|
def __exit__(self, *args):
|
|
self.close()
|