145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
"""Archive sniffing and hardened extraction (zip / 7z / rar)."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import stat
|
|
import subprocess
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from .util import safe_relpath
|
|
|
|
MAX_UNPACKED = 32 * 1024 ** 3
|
|
|
|
|
|
class ArchiveError(Exception):
|
|
pass
|
|
|
|
|
|
def sniff(path: Path) -> str:
|
|
"""Return 'zip', '7z', 'rar' based on magic bytes (extensions lie)."""
|
|
with open(path, "rb") as fh:
|
|
head = fh.read(8)
|
|
if head[:4] in (b"PK\x03\x04", b"PK\x05\x06"):
|
|
return "zip"
|
|
if head[:6] == b"7z\xbc\xaf\x27\x1c":
|
|
return "7z"
|
|
if head[:4] == b"Rar!":
|
|
return "rar"
|
|
raise ArchiveError(f"{Path(path).name}: not a zip, 7z or rar archive")
|
|
|
|
|
|
def _validate_name(name: str) -> str:
|
|
rel = safe_relpath(name)
|
|
if rel is None:
|
|
raise ArchiveError(f"unsafe path in archive: {name!r}")
|
|
return rel
|
|
|
|
|
|
def _extract_zip(path: Path, dest: Path) -> None:
|
|
try:
|
|
zf = zipfile.ZipFile(path)
|
|
except zipfile.BadZipFile as e:
|
|
raise ArchiveError(f"corrupt zip: {e}") from e
|
|
with zf:
|
|
infos = zf.infolist()
|
|
# validate everything before writing anything
|
|
total = 0
|
|
plan = []
|
|
for info in infos:
|
|
rel = _validate_name(info.filename)
|
|
mode = info.external_attr >> 16
|
|
if stat.S_ISLNK(mode):
|
|
raise ArchiveError(f"symlink entries are not allowed: {info.filename!r}")
|
|
total += info.file_size
|
|
plan.append((info, rel))
|
|
if total > MAX_UNPACKED:
|
|
raise ArchiveError("archive unpacks to an unreasonable size")
|
|
root = os.path.realpath(dest)
|
|
for info, rel in plan:
|
|
target = os.path.join(root, rel)
|
|
if os.path.commonpath([root, os.path.realpath(target)]) != root:
|
|
raise ArchiveError(f"unsafe path in archive: {info.filename!r}")
|
|
if info.is_dir():
|
|
os.makedirs(target, exist_ok=True)
|
|
continue
|
|
os.makedirs(os.path.dirname(target), exist_ok=True)
|
|
with zf.open(info) as src, open(target, "wb") as out:
|
|
shutil.copyfileobj(src, out)
|
|
|
|
|
|
def _tool(*names: str) -> str | None:
|
|
for n in names:
|
|
p = shutil.which(n)
|
|
if p:
|
|
return p
|
|
return None
|
|
|
|
|
|
def _run(cmd: list[str]) -> str:
|
|
try:
|
|
r = subprocess.run(cmd, capture_output=True, text=True, errors="replace", timeout=3600)
|
|
except FileNotFoundError as e:
|
|
raise ArchiveError(f"required tool missing: {cmd[0]}") from e
|
|
if r.returncode != 0:
|
|
raise ArchiveError((r.stderr or r.stdout).strip().splitlines()[-1:][0] if (r.stderr or r.stdout).strip()
|
|
else f"{cmd[0]} failed")
|
|
return r.stdout
|
|
|
|
|
|
def _list_7z(tool: str, path: Path) -> list[str]:
|
|
out = _run([tool, "l", "-slt", "-ba", "--", str(path)])
|
|
return [line[7:] for line in out.splitlines() if line.startswith("Path = ")]
|
|
|
|
|
|
def _forbid_symlinks(dest: Path) -> None:
|
|
for dirpath, dirnames, filenames in os.walk(dest):
|
|
for n in dirnames + filenames:
|
|
p = os.path.join(dirpath, n)
|
|
if os.path.islink(p):
|
|
raise ArchiveError(f"symlink found in extracted archive: {n!r}")
|
|
|
|
|
|
def _extract_7z(path: Path, dest: Path) -> None:
|
|
tool = _tool("7z", "7zz", "7za", "7zr")
|
|
if not tool:
|
|
raise ArchiveError("7z is not installed")
|
|
for n in _list_7z(tool, path):
|
|
_validate_name(n)
|
|
_run([tool, "x", "-y", f"-o{dest}", "--", str(path)])
|
|
_forbid_symlinks(dest)
|
|
|
|
|
|
def _extract_rar(path: Path, dest: Path) -> None:
|
|
unrar = _tool("unrar")
|
|
if unrar:
|
|
names = _run([unrar, "lb", "--", str(path)]).splitlines()
|
|
for n in names:
|
|
if n.strip():
|
|
_validate_name(n)
|
|
_run([unrar, "x", "-y", "-o+", "--", str(path), str(dest) + os.sep])
|
|
_forbid_symlinks(dest)
|
|
else:
|
|
_extract_7z(path, dest) # 7z can read rar when its codec is present
|
|
|
|
|
|
def extract(archive: Path, dest: Path) -> None:
|
|
"""Extract into an existing empty directory, rejecting traversal and symlinks."""
|
|
archive, dest = Path(archive), Path(dest)
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
kind = sniff(archive)
|
|
{"zip": _extract_zip, "7z": _extract_7z, "rar": _extract_rar}[kind](archive, dest)
|
|
|
|
|
|
def copy_tree(src: Path, dest: Path) -> None:
|
|
"""Import an already-extracted folder; symlinks are skipped."""
|
|
for dirpath, dirnames, filenames in os.walk(src):
|
|
rel = Path(dirpath).relative_to(src)
|
|
(dest / rel).mkdir(parents=True, exist_ok=True)
|
|
dirnames[:] = [d for d in dirnames if not os.path.islink(os.path.join(dirpath, d))]
|
|
for f in filenames:
|
|
s = os.path.join(dirpath, f)
|
|
if not os.path.islink(s):
|
|
shutil.copy2(s, dest / rel / f)
|