75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
"""Validated archive extraction with no links or traversal."""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import hmac
|
|
import io
|
|
import os
|
|
from pathlib import Path, PurePosixPath
|
|
import tarfile
|
|
|
|
try:
|
|
from tools.aios_spark_runner.errors import ConfigError
|
|
from tools.aios_spark_runner.securefs import (
|
|
atomic_write,
|
|
read_private_bytes,
|
|
reject_symlink_chain,
|
|
)
|
|
except ModuleNotFoundError:
|
|
from aios_spark_runner.errors import ConfigError
|
|
from aios_spark_runner.securefs import (
|
|
atomic_write,
|
|
read_private_bytes,
|
|
reject_symlink_chain,
|
|
)
|
|
|
|
|
|
def extract_artifact(
|
|
artifact: Path,
|
|
destination: Path,
|
|
*,
|
|
expected_digest: str,
|
|
expected_uid: int | None,
|
|
) -> None:
|
|
artifact = Path(artifact)
|
|
destination = Path(destination)
|
|
reject_symlink_chain(destination)
|
|
seen: set[str] = set()
|
|
artifact_bytes = read_private_bytes(artifact, expected_uid=expected_uid)
|
|
if not hmac.compare_digest(hashlib.sha256(artifact_bytes).hexdigest(), expected_digest):
|
|
raise ConfigError("artifact digest mismatch during extraction")
|
|
|
|
try:
|
|
archive = tarfile.open(fileobj=io.BytesIO(artifact_bytes), mode="r:gz")
|
|
except (tarfile.TarError, OSError) as exc:
|
|
raise ConfigError("runner artifact is not a valid tar.gz") from exc
|
|
with archive:
|
|
members = archive.getmembers()
|
|
for member in members:
|
|
pure = PurePosixPath(member.name)
|
|
if pure.is_absolute() or ".." in pure.parts or not pure.parts:
|
|
raise ConfigError("unsafe artifact member path")
|
|
if member.name in seen:
|
|
raise ConfigError("duplicate artifact member")
|
|
seen.add(member.name)
|
|
if not member.isfile():
|
|
raise ConfigError("artifact links and special files are forbidden")
|
|
for member in sorted(members, key=lambda item: item.name):
|
|
target = destination.joinpath(*PurePosixPath(member.name).parts)
|
|
target.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
reject_symlink_chain(target.parent)
|
|
stream = archive.extractfile(member)
|
|
if stream is None:
|
|
raise ConfigError("artifact member could not be read")
|
|
payload = stream.read()
|
|
if len(payload) != member.size:
|
|
raise ConfigError("artifact member size mismatch")
|
|
mode = 0o755 if member.name == "bin/aios-spark-runner" else 0o600
|
|
atomic_write(target, payload, mode=mode)
|
|
for directory in sorted(
|
|
(item for item in destination.rglob("*") if item.is_dir()),
|
|
key=lambda item: len(item.parts),
|
|
reverse=True,
|
|
):
|
|
os.chmod(directory, 0o700)
|