File size: 5,108 Bytes
ed3aeeb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
"""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())
|