Download Files with HTTP Range Resume

from __future__ import annotations

import http.server
import os
import tempfile
import threading
import urllib.request

PAYLOAD = b"0123456789" * 1000


class RangeHandler(http.server.BaseHTTPRequestHandler):
    def do_GET(self) -> None:
        header = self.headers.get("Range", "")
        start = 0
        if header.startswith("bytes="):
            start = int(header.removeprefix("bytes=").split("-", 1)[0] or 0)
        body = PAYLOAD[start:]
        self.send_response(206 if start else 200)
        self.send_header("Content-Length", str(len(body)))
        if start:
            self.send_header(
                "Content-Range",
                f"bytes {start}-{len(PAYLOAD) - 1}/{len(PAYLOAD)}",
            )
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format: str, *args: object) -> None:
        return


def download_with_resume(url: str, destination: str, chunk_size: int = 4096) -> int:
    existing = os.path.getsize(destination) if os.path.exists(destination) else 0
    request = urllib.request.Request(url)
    if existing:
        request.add_header("Range", f"bytes={existing}-")

    with urllib.request.urlopen(request, timeout=5) as response:
        if existing and response.status == 206:
            mode = "ab"
        else:
            mode = "wb"
            existing = 0
        with open(destination, mode) as handle:
            while True:
                chunk = response.read(chunk_size)
                if not chunk:
                    break
                handle.write(chunk)
    return os.path.getsize(destination)


def main() -> None:
    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), RangeHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        with tempfile.TemporaryDirectory() as raw_dir:
            destination = os.path.join(raw_dir, "payload.bin")
            with open(destination, "wb") as handle:
                handle.write(PAYLOAD[:5000])

            url = f"http://127.0.0.1:{server.server_address[1]}/payload.bin"
            size = download_with_resume(url, destination)
            with open(destination, "rb") as handle:
                downloaded = handle.read()
            print(f"downloaded {size} bytes")
            assert downloaded == PAYLOAD
    finally:
        server.shutdown()
        server.server_close()
        thread.join(timeout=5)


if __name__ == "__main__":
    main()