422 lines
14 KiB
Plaintext
422 lines
14 KiB
Plaintext
import string
|
|
import stdio
|
|
import stdlib
|
|
import requests
|
|
import requests.https as rhttps
|
|
import vthreading as vt
|
|
import win32.win32file as wf
|
|
import win32.win32base as wb
|
|
import win32.win32sync as wsync
|
|
|
|
|
|
# ============================================================
|
|
# requests/download.vp - large-file ranged download (Range + serial/multi-threaded)
|
|
#
|
|
# Usage:
|
|
# import requests.download as rd
|
|
# rc: int = rd.download("https://host/big.bin", "big.bin", threads=4,
|
|
# chunk_size=262144, verify=True, ca_file="ca.crt")
|
|
# if rc != 0:
|
|
# print(rd.last_error())
|
|
#
|
|
# Strategy:
|
|
# threads == 1 serial Range chunks (fetched and written to the file in order)
|
|
# threads > 1 multi-threaded: the file is split into bands by thread count, each thread uses one
|
|
# connection to download its band's chunks in order, writing at the corresponding offsets of the same file.
|
|
#
|
|
# Notes:
|
|
# - probes total size and server Range support via Content-Range;
|
|
# - http:// goes through requests (pure socket), https:// goes through requests.https (openssl);
|
|
# - each chunk request omits Accept-Encoding, the response is raw bytes and is written straight to disk.
|
|
# ============================================================
|
|
|
|
define DL_MAX_THREADS = 8
|
|
define DL_DEFAULT_CHUNK = 262144
|
|
define DL_REDIRECT_MAX = 5
|
|
|
|
_dl_err_store: [512]char
|
|
|
|
|
|
def last_error() -> str:
|
|
"""Last download failure reason (empty string if none)."""
|
|
return _dl_err_store
|
|
|
|
|
|
def _dl_set_err(text: str):
|
|
if text is None:
|
|
return
|
|
_dl_err_store[0] = 0
|
|
stdio.snprintf(_dl_err_store, 512, "%s", text)
|
|
|
|
|
|
def _dl_strdup(s: str) -> str:
|
|
"""Duplicate string on stdlib heap (caller frees); None on failure."""
|
|
if s is None:
|
|
return None
|
|
n: sizet = string.strlen(s)
|
|
out: str = stdlib.malloc(n + 1)
|
|
if out is None:
|
|
return None
|
|
if n > 0:
|
|
string.memcpy(out, s, n)
|
|
out[n] = 0
|
|
return out
|
|
|
|
|
|
# ============================================================
|
|
# Single-slice fetch (redirect follow, <=5 hops; dispatch by http/https scheme)
|
|
# ============================================================
|
|
|
|
def _dl_slice(url: str, pos: sizet, want: int, verify: bool, ca_file: str,
|
|
dst: str, dst_cap: int, out_total: sizet*) -> int:
|
|
"""Fetch slice [pos, pos+want); body bytes written to dst (locate body via _dl_body_off).
|
|
Returns -2 on too many redirects; -3 on network/server failure.
|
|
"""
|
|
if url is None or dst is None or out_total is None:
|
|
return -1
|
|
cur: str = _dl_strdup(url)
|
|
if cur is None:
|
|
return -1
|
|
hops: int = 0
|
|
rc: int = 0
|
|
while True:
|
|
is_https: int = 1 if (string.strncmp(cur, "https://", 8) == 0) else 0
|
|
st: int = 0
|
|
tot: sizet = 0
|
|
loc: [512]char
|
|
if is_https != 0:
|
|
rc = rhttps._DlOnce(cur, pos, want, verify, ca_file, dst, dst_cap,
|
|
&st, loc, 512, &tot)
|
|
else:
|
|
rc = requests._DlOnce(cur, pos, want, dst, dst_cap,
|
|
&st, loc, 512, &tot)
|
|
if rc == -2 and loc[0] != 0:
|
|
hops += 1
|
|
if hops > DL_REDIRECT_MAX:
|
|
_dl_set_err("too many redirects")
|
|
stdlib.free(cur)
|
|
return -2
|
|
new_url: str = stdlib.malloc(2048)
|
|
if new_url is None:
|
|
stdlib.free(cur)
|
|
return -1
|
|
if requests._ResolveRedirectUrl(cur, loc, new_url, 2048) != 0:
|
|
stdlib.free(new_url)
|
|
stdlib.free(cur)
|
|
_dl_set_err("bad redirect location")
|
|
return -3
|
|
stdlib.free(cur)
|
|
cur = new_url
|
|
continue
|
|
break
|
|
stdlib.free(cur)
|
|
if rc < 0:
|
|
return rc
|
|
*out_total = tot
|
|
return rc
|
|
|
|
|
|
def _dl_body_off(dst: str, dst_cap: int) -> int:
|
|
"""Offset of response body (after headers) in dst; -1 if not found."""
|
|
return requests._FindBodyStart(dst, dst_cap)
|
|
|
|
|
|
# ============================================================
|
|
# File write (write at offset; supports concurrent threads writing distinct ranges)
|
|
# ============================================================
|
|
|
|
def _dl_write_at(path: str, src: str, length: int, offset: sizet) -> int:
|
|
"""Open file (create if absent), seek to offset, write length bytes. 0 ok; -1 fail."""
|
|
if path is None or src is None or length <= 0:
|
|
return -1
|
|
h: wb.HANDLE = wf.CreateFileA(
|
|
path, wf.GENERIC_WRITE,
|
|
wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
|
|
None, wf.OPEN_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
|
|
if h == wb.INVALID_HANDLE_VALUE:
|
|
return -1
|
|
# Use 32-bit SetFilePointer (lDistanceToMove passed by value as long; no struct ABI ambiguity).
|
|
# lpDistanceToMoveHigh=NULL limits distance to 32-bit signed (<2GiB), so guard the upper bound first.
|
|
if offset > sizet(0x7FFFFFFF):
|
|
wb.CloseHandle(h)
|
|
return -1
|
|
ok_seek: long = wf.SetFilePointer(h, long(offset), None, wf.FILE_BEGIN)
|
|
# INVALID_SET_FILE_POintER(0xFFFFFFFF) truncates to long -1; offset is already clamped to
|
|
# 0..0x7FFFFFFF, so 0xFFFFFFFF cannot be a valid offset and no GetLastError check is needed.
|
|
if ok_seek == -1:
|
|
wb.CloseHandle(h)
|
|
return -1
|
|
written: ulong = 0
|
|
ok_w: int = wf.WriteFile(h, src, uint32(length), &written, None)
|
|
wb.CloseHandle(h)
|
|
if ok_w == 0:
|
|
return -1
|
|
if int(written) != length:
|
|
return -1
|
|
return 0
|
|
|
|
|
|
def _dl_truncate(path: str) -> int:
|
|
"""Create (truncate) an empty file via CREATE_ALWAYS. 0 ok; -1 fail."""
|
|
h: wb.HANDLE = wf.CreateFileA(
|
|
path, wf.GENERIC_WRITE,
|
|
wf.FILE_SHARE_READ | wf.FILE_SHARE_WRITE | wf.FILE_SHARE_DELETE,
|
|
None, wf.CREATE_ALWAYS, wf.FILE_ATTRIBUTE_NORMAL, None)
|
|
if h == wb.INVALID_HANDLE_VALUE:
|
|
return -1
|
|
wb.CloseHandle(h)
|
|
return 0
|
|
|
|
|
|
# ============================================================
|
|
# Probe: total size / Range support (first slice bytes=0-0)
|
|
# ============================================================
|
|
|
|
def _dl_probe(url: str, verify: bool, ca_file: str,
|
|
out_total: sizet*) -> int:
|
|
"""Probe server total size and Range support. Returns 0=Range supported (out_total valid);
|
|
1=server returned 200 (no Range, full download, out_total=Content-Length); -3 fail.
|
|
"""
|
|
*out_total = 0
|
|
buf: str = stdlib.malloc(8192)
|
|
if buf is None:
|
|
_dl_set_err("out of memory")
|
|
return -3
|
|
tot: sizet = 0
|
|
rc: int = _dl_slice(url, 0, 0, verify, ca_file, buf, 8192, &tot)
|
|
stdlib.free(buf)
|
|
if rc < 0:
|
|
_dl_set_err("probe failed")
|
|
return -3
|
|
if tot > 0:
|
|
*out_total = tot
|
|
return 0
|
|
_dl_set_err("cannot determine size")
|
|
return -3
|
|
|
|
|
|
# ============================================================
|
|
# Download one band ([b_start, b_end)), fetching and writing slice by slice
|
|
# ============================================================
|
|
|
|
def _dl_download_band(path: str, url: str, is_https: int, verify: bool, ca_file: str,
|
|
b_start: sizet, b_end: sizet, chunk: int,
|
|
buf: str, buf_cap: int,
|
|
p_done: sizet*) -> int:
|
|
"""0 ok; -1 network/server; -2 write failure. p_done accumulates bytes written."""
|
|
pos: sizet = b_start
|
|
while pos < b_end:
|
|
want: int = chunk
|
|
if b_end - pos < sizet(want):
|
|
want = int(b_end - pos)
|
|
tot: sizet = 0
|
|
rc: int = _dl_slice(url, pos, want, verify, ca_file, buf, buf_cap, &tot)
|
|
if rc <= 0:
|
|
_dl_set_err("slice fetch failed")
|
|
return -1
|
|
bo: int = _dl_body_off(buf, buf_cap)
|
|
if bo < 0:
|
|
_dl_set_err("bad slice response")
|
|
return -1
|
|
got: int = rc
|
|
if bo + got > buf_cap:
|
|
got = buf_cap - bo
|
|
if got <= 0:
|
|
_dl_set_err("slice empty")
|
|
return -1
|
|
if _dl_write_at(path, buf + bo, got, pos) != 0:
|
|
_dl_set_err("write file failed")
|
|
return -2
|
|
*p_done = *p_done + sizet(got)
|
|
pos += sizet(got)
|
|
if got < want:
|
|
# Server ended early (only the last slice is < want on normal completion)
|
|
if pos >= b_end:
|
|
break
|
|
_dl_set_err("short slice")
|
|
return -1
|
|
return 0
|
|
|
|
|
|
# ============================================================
|
|
# Multi-thread worker
|
|
# ============================================================
|
|
class DlJob:
|
|
path: str
|
|
url: str
|
|
is_https: int
|
|
verify: int
|
|
ca_file: str
|
|
start: sizet
|
|
end: sizet
|
|
chunk: int
|
|
buf: str
|
|
buf_cap: int
|
|
done: sizet
|
|
status: int
|
|
|
|
|
|
def _dl_worker(arg: rawptr) -> rawptr:
|
|
if arg is None:
|
|
return None
|
|
job: DlJob* = arg
|
|
rc: int = _dl_download_band(
|
|
job.path, job.url, job.is_https,
|
|
True if job.verify != 0 else False,
|
|
job.ca_file,
|
|
job.start, job.end, job.chunk, job.buf, job.buf_cap,
|
|
&job.done)
|
|
if rc != 0:
|
|
job.status = rc
|
|
return None
|
|
|
|
|
|
# ============================================================
|
|
# Public API
|
|
# ============================================================
|
|
|
|
def download(url: str, path: str, threads: int = 1,
|
|
chunk_size: int = DL_DEFAULT_CHUNK,
|
|
verify: bool = True, ca_file: str = None) -> int:
|
|
"""Chunked large-file download over http:// or https:// (Range). threads 1=serial,
|
|
2~8=parallel bands; chunk_size bytes (default 256KB). Returns 0/-1/-2/-3/-4 (see last_error()).
|
|
"""
|
|
_dl_err_store[0] = 0
|
|
if url is None or path is None or url[0] == 0 or path[0] == 0:
|
|
_dl_set_err("url/path empty")
|
|
return -1
|
|
if chunk_size < 1024:
|
|
chunk_size = DL_DEFAULT_CHUNK
|
|
if threads < 1:
|
|
threads = 1
|
|
if threads > DL_MAX_THREADS:
|
|
threads = DL_MAX_THREADS
|
|
if string.strncmp(url, "http://", 7) != 0 \
|
|
and string.strncmp(url, "https://", 8) != 0:
|
|
_dl_set_err("unsupported scheme")
|
|
return -1
|
|
|
|
# Probe total size and Range support
|
|
total: sizet = 0
|
|
prc: int = _dl_probe(url, verify, ca_file, &total)
|
|
if prc < 0:
|
|
return -3
|
|
if total == 0:
|
|
_dl_set_err("cannot determine size")
|
|
return -3
|
|
|
|
# Create (truncate) target file
|
|
if _dl_truncate(path) != 0:
|
|
_dl_set_err("create file failed")
|
|
return -2
|
|
|
|
stdio.printf("[download] %s -> %s total=%zu bytes threads=%d chunk=%d\n",
|
|
url, path, total, threads, chunk_size)
|
|
|
|
slab: int = chunk_size + 8192
|
|
|
|
if threads == 1:
|
|
# ---- Serial slices ----
|
|
buf: str = stdlib.malloc(sizet(slab))
|
|
if buf is None:
|
|
_dl_set_err("out of memory")
|
|
return -3
|
|
done: sizet = 0
|
|
rc: int = _dl_download_band(path, url, 0, verify, ca_file,
|
|
0, total, chunk_size, buf, slab, &done)
|
|
stdlib.free(buf)
|
|
if rc != 0:
|
|
return -3 if rc == -1 else -4
|
|
stdio.printf("[download] done: %zu/%zu bytes\n", done, total)
|
|
return 0
|
|
|
|
# ---- Multi-thread parallel bands ----
|
|
band: sizet = (total + sizet(threads) - 1) / sizet(threads)
|
|
ctxs: [DL_MAX_THREADS]rawptr
|
|
handles: [DL_MAX_THREADS]rawptr
|
|
threads_ok: int = 0
|
|
|
|
for i in range(int(threads)):
|
|
b_start: sizet = sizet(i) * band
|
|
if b_start >= total:
|
|
break
|
|
b_end: sizet = b_start + band
|
|
if b_end > total:
|
|
b_end = total
|
|
job: DlJob* = stdlib.calloc(1, DlJob.__sizeof__())
|
|
if job is None:
|
|
break
|
|
job.path = _dl_strdup(path)
|
|
job.url = _dl_strdup(url)
|
|
if job.path is None or job.url is None:
|
|
stdlib.free(job.path)
|
|
stdlib.free(job.url)
|
|
stdlib.free(rawptr(job))
|
|
break
|
|
job.is_https = 1 if (string.strncmp(url, "https://", 8) == 0) else 0
|
|
job.verify = 1 if verify else 0
|
|
job.ca_file = _dl_strdup(ca_file) if ca_file is not None else None
|
|
job.start = b_start
|
|
job.end = b_end
|
|
job.chunk = chunk_size
|
|
job.buf_cap = slab
|
|
job.buf = stdlib.malloc(sizet(slab))
|
|
if job.buf is None:
|
|
stdlib.free(job.path)
|
|
stdlib.free(job.url)
|
|
if job.ca_file is not None:
|
|
stdlib.free(job.ca_file)
|
|
stdlib.free(rawptr(job))
|
|
break
|
|
job.done = 0
|
|
job.status = 0
|
|
ctxs[i] = rawptr(job)
|
|
|
|
th: vt.Thread* = vt.Thread(_dl_worker, ctxs[i])
|
|
if th is None:
|
|
stdlib.free(job.buf)
|
|
stdlib.free(job.path)
|
|
stdlib.free(job.url)
|
|
if job.ca_file is not None:
|
|
stdlib.free(job.ca_file)
|
|
stdlib.free(rawptr(job))
|
|
break
|
|
th.start()
|
|
handles[i] = th.handle
|
|
threads_ok += 1
|
|
|
|
if threads_ok == 0:
|
|
_dl_set_err("thread create failed")
|
|
return -3
|
|
|
|
# Wait for all threads to finish
|
|
for j in range(int(threads_ok)):
|
|
wsync.WaitForSingleObject(handles[j], 0xFFFFFFFF)
|
|
wb.CloseHandle(handles[j])
|
|
|
|
# Aggregate results
|
|
total_done: sizet = 0
|
|
failed_job: int = 0
|
|
for k in range(int(threads_ok)):
|
|
jobk: DlJob* = ctxs[k]
|
|
total_done += jobk.done
|
|
if jobk.status != 0:
|
|
failed_job = 1
|
|
|
|
# Free thread contexts
|
|
for m in range(int(threads_ok)):
|
|
jobm: DlJob* = ctxs[m]
|
|
stdlib.free(jobm.path)
|
|
stdlib.free(jobm.url)
|
|
if jobm.ca_file is not None:
|
|
stdlib.free(jobm.ca_file)
|
|
stdlib.free(jobm.buf)
|
|
stdlib.free(rawptr(jobm))
|
|
|
|
if failed_job != 0 or total_done != total:
|
|
stdio.printf("[download] incomplete: %zu/%zu bytes\n", total_done, total)
|
|
_dl_set_err("download incomplete")
|
|
return -3
|
|
stdio.printf("[download] done: %zu/%zu bytes\n", total_done, total)
|
|
return 0
|