| |
| """Resume a public immutable artifact with independent HTTP byte ranges.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import concurrent.futures |
| import hashlib |
| import json |
| import os |
| import shutil |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
|
|
| def digest(path: Path, algorithm: str) -> str: |
| value = hashlib.new(algorithm) |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| value.update(chunk) |
| return value.hexdigest() |
|
|
|
|
| def fetch_part(url: str, path: Path, start: int, end: int, retries: int) -> dict[str, object]: |
| expected = end - start + 1 |
| path.parent.mkdir(parents=True, exist_ok=True) |
| for attempt in range(1, retries + 1): |
| current = path.stat().st_size if path.exists() else 0 |
| if current == expected: |
| return {"path": str(path), "bytes": current, "reused": True, "attempts": attempt - 1} |
| if current > expected: |
| raise RuntimeError(f"oversized range file {path}: {current} > {expected}") |
| request_start = start + current |
| request = urllib.request.Request( |
| url, |
| headers={ |
| "Range": f"bytes={request_start}-{end}", |
| "User-Agent": "mcu-ir-research/1.0", |
| }, |
| ) |
| try: |
| with urllib.request.urlopen(request, timeout=120) as response: |
| status = getattr(response, "status", None) |
| content_range = response.headers.get("Content-Range", "") |
| if status != 206 or not content_range.startswith(f"bytes {request_start}-{end}/"): |
| raise RuntimeError(f"unexpected range response: status={status} content-range={content_range!r}") |
| with path.open("ab") as output: |
| shutil.copyfileobj(response, output, length=1024 * 1024) |
| except Exception: |
| if attempt == retries: |
| raise |
| time.sleep(min(2**attempt, 15)) |
| raise AssertionError("unreachable") |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--url", required=True) |
| parser.add_argument("--output", required=True, type=Path) |
| parser.add_argument("--expected-bytes", required=True, type=int) |
| parser.add_argument("--expected-md5", required=True) |
| parser.add_argument("--parts", type=int, default=16) |
| parser.add_argument("--retries", type=int, default=5) |
| args = parser.parse_args() |
| output = args.output.resolve() |
| if output.exists(): |
| actual_md5 = digest(output, "md5") |
| if output.stat().st_size == args.expected_bytes and actual_md5 == args.expected_md5: |
| print(json.dumps({"status": "PASS", "reused": True, "output": str(output), "bytes": output.stat().st_size, "md5": actual_md5, "sha256": digest(output, "sha256")}, sort_keys=True)) |
| return 0 |
| raise SystemExit(f"refusing to overwrite invalid existing output: {output}") |
|
|
| part_dir = output.parent / f".{output.name}.parts" |
| part_dir.mkdir(parents=True, exist_ok=True) |
| base, remainder = divmod(args.expected_bytes, args.parts) |
| ranges: list[tuple[int, Path, int, int]] = [] |
| offset = 0 |
| for index in range(args.parts): |
| size = base + (1 if index < remainder else 0) |
| ranges.append((index, part_dir / f"part-{index:03d}", offset, offset + size - 1)) |
| offset += size |
| with concurrent.futures.ThreadPoolExecutor(max_workers=args.parts) as executor: |
| futures = { |
| executor.submit(fetch_part, args.url, path, start, end, args.retries): index |
| for index, path, start, end in ranges |
| } |
| for future in concurrent.futures.as_completed(futures): |
| result = future.result() |
| print(json.dumps({"part": futures[future], **result}, sort_keys=True), flush=True) |
|
|
| temporary = output.with_suffix(output.suffix + ".assembling") |
| md5 = hashlib.md5(usedforsecurity=False) |
| sha256 = hashlib.sha256() |
| with temporary.open("wb") as destination: |
| for _, path, start, end in ranges: |
| expected = end - start + 1 |
| if path.stat().st_size != expected: |
| raise RuntimeError(f"range size mismatch for {path}") |
| with path.open("rb") as source: |
| for chunk in iter(lambda: source.read(1024 * 1024), b""): |
| destination.write(chunk) |
| md5.update(chunk) |
| sha256.update(chunk) |
| destination.flush() |
| os.fsync(destination.fileno()) |
| if temporary.stat().st_size != args.expected_bytes or md5.hexdigest() != args.expected_md5: |
| raise RuntimeError( |
| f"assembled artifact mismatch: bytes={temporary.stat().st_size} md5={md5.hexdigest()}" |
| ) |
| os.replace(temporary, output) |
| print(json.dumps({"status": "PASS", "reused": False, "output": str(output), "bytes": output.stat().st_size, "md5": md5.hexdigest(), "sha256": sha256.hexdigest()}, sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
