acquire_inputs.py

1"""One-stream, bounded acquisition of the two books and complete small tables."""
2import argparse, base64, hashlib, json, os, re, shutil, time, urllib.request, zipfile
3from pathlib import Path
4from curriculum_policy import preflight, TASK_CAP, FREE_RESERVE
5
6ROOT = Path('/work/littleply-curriculum-20261029')
7LOCAL = Path(__file__).resolve().parent
8MIRROR = 'https://tablebase.lichess.ovh/tables/standard/'
9BOOKS = ['8moves_v3.pgn', 'bjbraams_chessdb_198350_lines.pgn']
10
11def sha(path):
12    h=hashlib.sha256()
13    with Path(path).open('rb') as f:
14        for block in iter(lambda:f.read(1<<20), b''): h.update(block)
15    return h.hexdigest()
16
17def allocation():
18    return sum(p.stat().st_size for root in (ROOT,LOCAL) for p in root.rglob('*') if p.is_file())
19
20def guard(extra=0):
21    if allocation()+extra > TASK_CAP: raise RuntimeError('2GB task cap')
22    if min(shutil.disk_usage(ROOT).free,shutil.disk_usage('/mnt/c').free)-extra < FREE_RESERVE:
23        raise RuntimeError('50GB reserve on both filesystems')
24
25def save(name,value):
26    data=json.dumps(value,indent=2).encode(); guard(len(data))
27    path=ROOT/name; temp=path.with_suffix(path.suffix+'.tmp')
28    temp.write_bytes(data); temp.replace(path)
29
30def small(url,name):
31    p=ROOT/name
32    if p.exists(): return p.read_bytes()
33    guard(2_000_000)
34    with urllib.request.urlopen(url,timeout=30) as r: data=r.read(2_000_001)
35    if len(data)>2_000_000: raise ValueError('metadata response oversized')
36    p.write_bytes(data)
37    return data
38
39def parse_listing(data,suffix):
40    found={}
41    for line in data.decode().splitlines():
42        match=re.search(r'<a href="(K[QRBNP]*vK[QRBNP]*\.'+suffix+r')">.*</a>\s+\S+\s+\S+\s+(\d+)\s*$',line)
43        if match:
44            name,size=match.groups()
45            if name in found: raise ValueError('duplicate listing')
46            found[name]=int(size)
47    if not found: raise ValueError('empty listing')
48    return found
49
50def parse_hashes(data):
51    found={}
52    for line in data.decode().splitlines():
53        h,n=line.split()
54        if not re.fullmatch('[0-9a-f]{64}',h): raise ValueError('bad SHA256')
55        if n in found and found[n]!=h: raise ValueError('conflicting SHA256')
56        found[n]=h
57    return found
58
59def book_sri(content):
60    normalized=content.replace(b'\r\n',b'\n').replace(b'\r',b'\n')
61    return base64.b64encode(hashlib.sha384(normalized).digest()).decode()
62
63def metadata():
64    ROOT.mkdir(exist_ok=True)
65    guard()
66    before={'host_free':shutil.disk_usage('/mnt/c').free,'native_free':shutil.disk_usage(ROOT).free,'task_bytes':allocation()}
67    revisions={}
68    for repo in ('official-stockfish/books','syzygy1/tb'):
69        revisions[repo]=json.loads(small('https://api.github.com/repos/'+repo+'/commits/master',repo.replace('/','-')+'-commit.json'))['sha']
70    br='https://raw.githubusercontent.com/official-stockfish/books/'+revisions['official-stockfish/books']+'/'
71    sr='https://raw.githubusercontent.com/syzygy1/tb/'+revisions['syzygy1/tb']+'/'
72    books=json.loads(small(br+'books.json','books.json'))
73    small(br+'LICENSE','books-LICENSE.txt')
74    small(br+'update_json.py','upstream-update_json.txt')
75    small(sr+'README.md','syzygy-README.txt')
76    small(sr+'src/checksum.c','upstream-checksum.txt')
77    hashes=parse_hashes(small(MIRROR+'sha256','mirror-sha256.txt'))
78    rows=[]
79    for kind,suffix in (('wdl','rtbw'),('dtz','rtbz')):
80        known=small(sr+'checksums/'+kind+'345.txt',kind+'345.txt').decode()
81        known=dict(line.split(': ') for line in known.splitlines())
82        listing=parse_listing(small(MIRROR+'3-4-5-'+kind+'/',kind+'-listing.html'),suffix)
83        if set(known)!=set(listing): raise ValueError('upstream/mirror filename set mismatch')
84        for name,size in sorted(listing.items()):
85            rows.append(dict(name=name,bytes=size,url=MIRROR+'3-4-5-'+kind+'/'+name,
86                             sha256=hashes[name],upstream_embedded_checksum=known[name],kind='table'))
87    reusable={}
88    for r in rows:
89        p=ROOT/r['name']
90        if p.exists():
91            if p.stat().st_size!=r['bytes'] or sha(p)!=r['sha256']: raise ValueError('existing table mismatch')
92            reusable[r['name']]=r['bytes']
93    check=preflight(rows,host_free=before['host_free'],native_free=before['native_free'],
94                    task_bytes=before['task_bytes'],reusable=reusable)
95    for name in BOOKS:
96        url=br+name+'.zip'
97        with urllib.request.urlopen(urllib.request.Request(url,method='HEAD'),timeout=30) as r:
98            size=int(r.headers['Content-Length'])
99        if not 0<size<25_000_000: raise ValueError('unexpected book size')
100        rows.append(dict(name=name+'.zip',bytes=size,url=url,kind='book',manifest=books[name]))
101    guard(sum(r['bytes'] for r in rows if not (ROOT/r['name']).exists())+100_000)
102    plan=dict(revisions=revisions,before=before,preflight=check,files=rows,
103              expected_table_bytes=sum(r['bytes'] for r in rows if r['kind']=='table'),
104              expected_book_bytes=sum(r['bytes'] for r in rows if r['kind']=='book'))
105    save('download-plan.json',plan)
106    print(json.dumps({k:v for k,v in plan.items() if k!='files'}),flush=True)
107
108def stream_file(response,path,expected_size,expected_sha=None,rate=8_000_000,check=None):
109    """No resume: retain partial as evidence; no overwriting an existing partial."""
110    partial=path.with_suffix(path.suffix+'.part')
111    h=hashlib.sha256(); n=0; start=time.monotonic()
112    with partial.open('xb') as f:
113        while True:
114            chunk=response.read(min(1<<20,expected_size-n+1))
115            if not chunk: break
116            if n+len(chunk)>expected_size: raise ValueError('response exceeds approved size')
117            if check: check(len(chunk))
118            f.write(chunk); h.update(chunk); n+=len(chunk)
119            delay=n/rate-(time.monotonic()-start)
120            if delay>0: time.sleep(min(delay,0.2))
121        f.flush(); os.fsync(f.fileno())
122    if n!=expected_size: raise ValueError('truncated response')
123    digest=h.hexdigest()
124    if expected_sha and digest!=expected_sha: raise ValueError('SHA256 mismatch')
125    partial.replace(path)
126    return digest
127
128def download():
129    plan=json.loads((ROOT/'download-plan.json').read_text())
130    # Re-evaluate whole plan and disk limits immediately before payload begins.
131    rows=plan['files']; tables=[r for r in rows if r['kind']=='table']
132    reused={}
133    for r in tables:
134        p=ROOT/r['name']
135        if p.exists():
136            if p.stat().st_size!=r['bytes'] or sha(p)!=r['sha256']: raise ValueError('reuse integrity')
137            reused[r['name']]=r['bytes']
138    preflight(tables,host_free=shutil.disk_usage('/mnt/c').free,native_free=shutil.disk_usage(ROOT).free,
139              task_bytes=allocation(),reusable=reused)
140    guard(sum(r['bytes'] for r in rows if not (ROOT/r['name']).exists())+100_000)
141    receipt=[]
142    for i,r in enumerate(rows):
143        p=ROOT/r['name']; reuse=p.exists()
144        if not reuse:
145            with urllib.request.urlopen(urllib.request.Request(r['url'],headers={'Accept-Encoding':'identity'}),timeout=60) as response:
146                if int(response.headers.get('Content-Length',-1))!=r['bytes']: raise ValueError('HTTP size changed')
147                stream_file(response,p,r['bytes'],r.get('sha256'),check=guard)
148        if p.stat().st_size!=r['bytes']: raise ValueError('local size mismatch')
149        digest=sha(p)
150        if r.get('sha256') and digest!=r['sha256']: raise ValueError('local hash mismatch')
151        entry={**r,'local_sha256':digest,'reused':reuse}
152        if r['kind']=='table':
153            if p.stat().st_size % 64 != 16: raise ValueError('no embedded checksum')
154            with p.open('rb') as f:
155                f.seek(-16,2); embedded=f.read(16).hex()
156            if embedded!=r['upstream_embedded_checksum']: raise ValueError('upstream embedded checksum mismatch')
157            entry['upstream_embedded_checksum_match']=True
158        if r['kind']=='book':
159            with zipfile.ZipFile(p) as z:
160                member=r['name'][:-4]
161                if z.namelist()!=[member]: raise ValueError('unexpected archive members')
162                if z.getinfo(member).file_size>100_000_000: raise ValueError('oversized expanded book')
163                sri=book_sri(z.read(member))
164                if sri!=r['manifest']['sri']: raise ValueError('uncompressed SHA384 SRI mismatch')
165                entry.update(sri_verified=True,uncompressed_bytes=z.getinfo(member).file_size)
166        receipt.append(entry)
167        if i%20==0 or i==len(rows)-1:
168            save('download-progress.json',dict(complete=len(receipt),total=len(rows),bytes=sum(x['bytes'] for x in receipt),last=p.name))
169            print(json.dumps(dict(complete=len(receipt),total=len(rows),last=p.name)),flush=True)
170    save('download-receipt.json',dict(files=receipt,table_bytes=plan['expected_table_bytes'],
171        book_bytes=plan['expected_book_bytes'],task_bytes=allocation(),before=plan['before'],
172        after={'host_free':shutil.disk_usage('/mnt/c').free,'native_free':shutil.disk_usage(ROOT).free},
173        complete=True,table_hash_verification='full file SHA256 matches mirror HTTPS manifest',
174        upstream_embedded_verification='all footers match pinned syzygy1/tb values; payload integrity independently verified by full-file mirror SHA256'))
175
176if __name__=='__main__':
177    p=argparse.ArgumentParser();p.add_argument('phase',choices=['metadata','download']);args=p.parse_args()
178    globals()[args.phase]()
179