Extract Tar Archives Without Path Traversal

from __future__ import annotations

import io
import tarfile
import tempfile
from pathlib import Path, PurePosixPath


def safe_members(archive: tarfile.TarFile, destination: Path) -> list[tarfile.TarInfo]:
    root = destination.resolve()
    accepted: list[tarfile.TarInfo] = []
    for member in archive.getmembers():
        path = PurePosixPath(member.name)
        if path.is_absolute() or ".." in path.parts:
            raise ValueError(f"unsafe archive path: {member.name}")
        if member.issym() or member.islnk():
            raise ValueError(f"links are not allowed: {member.name}")
        if not (member.isfile() or member.isdir()):
            raise ValueError(f"unsupported member type: {member.name}")
        target = (root / path).resolve()
        if root not in target.parents and target != root:
            raise ValueError(f"member escapes destination: {member.name}")
        accepted.append(member)
    return accepted


def extract_archive(data: bytes, destination: Path) -> list[str]:
    destination.mkdir(parents=True, exist_ok=True)
    with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as archive:
        members = safe_members(archive, destination)
        archive.extractall(destination, members=members)
    return sorted(
        path.relative_to(destination).as_posix()
        for path in destination.rglob("*")
        if path.is_file()
    )


def build_archive(entries: list[tuple[str, bytes]]) -> bytes:
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w:gz") as archive:
        for name, payload in entries:
            info = tarfile.TarInfo(name)
            info.size = len(payload)
            archive.addfile(info, io.BytesIO(payload))
    return buffer.getvalue()


def main() -> None:
    with tempfile.TemporaryDirectory() as raw_root:
        root = Path(raw_root)
        destination = root / "extract"
        archive = build_archive([("a.txt", b"alpha"), ("nested/b.txt", b"beta")])
        files = extract_archive(archive, destination)
        print(files)
        assert files == ["a.txt", "nested/b.txt"]
        assert (destination / "nested" / "b.txt").read_bytes() == b"beta"

        malicious = build_archive([("../evil.txt", b"owned")])
        try:
            extract_archive(malicious, root / "unsafe")
        except ValueError as exc:
            print(f"rejected: {exc}")
        else:
            raise AssertionError("traversal archive was accepted")
        assert not (root / "evil.txt").exists()


if __name__ == "__main__":
    main()