"""Session 38: diagnose and guard OutGuess 0.4's invalid iterator state.
Usage: python outguess_crash_audit.py EXTRACTED_SOURCE IMAGE_DIR SESSION37_RESULTS OUTPUT
Requires gcc, make, Pillow. Source tarball URL/hash are in outguess_source.json.
Builds outside the image directory; never rewrites source JPEGs.
This product includes software developed by Niels Provos.
"""
from pathlib import Path
import collections,difflib,hashlib,json,os,random,re,shutil,subprocess,sys
from PIL import Image

def sha(b):return hashlib.sha256(b).hexdigest()
def run(cmd,cwd=None):
 r=subprocess.run(list(map(str,cmd)),cwd=cwd,capture_output=True,timeout=120)
 return {'returncode':r.returncode,'stdout':r.stdout.decode('utf8','replace'),'stderr':r.stderr.decode('utf8','replace')}
def extract(exe,p,dst,key=None):
 if dst.exists():dst.unlink()
 r=run([exe,*([] if key is None else ['-k',key]),'-r',p,dst]);b=dst.read_bytes() if dst.exists() else b''
 r.update(bytes=len(b),sha256=sha(b));return r

def summarize(results):
 rows=results['files'];shortfalls=[]
 for row in rows:
  log=row['guarded']['stderr']
  if 'skipmod=0' in log:
   remaining=int(re.search(r'AUDIT_REMAINING bytes=(\d+)',log).group(1))
   offset,bits=map(int,re.search(r'offset=(\d+) usable_bits=(\d+)',log).groups())
   shortfalls.append({'file':row['file'],'remaining_bytes':remaining,'remaining_bits':bits-offset,'required_bits':8*remaining,'shortfall':8*remaining-(bits-offset)})
 assert len(shortfalls)==41 and all(x['shortfall']>0 for x in shortfalls)
 return {'guarded_return_codes':dict(collections.Counter(str(x['guarded']['returncode']) for x in rows)),'lp2_guarded_return_codes':dict(collections.Counter(str(x['guarded']['returncode']) for x in rows[17:])),'lp2_raw_headers':dict(collections.Counter(re.search(r'AUDIT_RAW_HEADER (\w+)',x['guarded']['stderr']).group(1) for x in rows[17:])),'all_previously_successful_outputs_unchanged':all(x['baseline']['sha256']==x['guarded']['sha256'] for x in rows if x['previous_status']==0),'all_41_crashes_now_explicit_errors':all(x['guarded']['returncode']==2 for x in rows if x['previous_status']==-8),'shortfalls':shortfalls}

def main():
 source,images,old,out=[Path(p).resolve() for p in sys.argv[1:]];out.mkdir(parents=True,exist_ok=True)
 build=out/'build';shutil.copytree(source,build,dirs_exist_ok=True)
 (build/'config.h').write_text('#define HAVE_MMAP 1\n#define HAVE_SNPRINTF 1\n')
 shutil.copyfile(build/'src/jpeg-6b-steg/jconfig.cfg',build/'src/jpeg-6b-steg/jconfig.h')
 lib=run(['make','-C','src/jpeg-6b-steg','-f','makefile.ansi','CC=gcc','CFLAGS=-O0 -g -DHAVE_STDC_HEADERS','libjpeg.a'],build);assert lib['returncode']==0,lib
 flags=['gcc','-O0','-g','-fsanitize=integer-divide-by-zero','-fno-sanitize-recover=integer-divide-by-zero','-I.','-Isrc']
 sources=['src/'+x+'.c' for x in ['outguess','golay','arc','pnm','jpg','iterator','md5']]
 def compile_to(dest):
  r=run([*flags,*sources,'src/jpeg-6b-steg/libjpeg.a','-lm','-o',dest],build);assert r['returncode']==0,r;return r
 baseline=out/'outguess-original';compile_to(baseline)
 diagnosis=extract(baseline,images/'18.jpg',out/'diagnostic18.bin')
 assert 'division by zero' in diagnosis['stderr']
 # Guard the existing algorithm. Never clamp zero to one or manufacture payload bits.
 edits={}
 p=build/'src/iterator.c';a=p.read_text();b=a.replace('#include <sys/types.h>','#include <sys/types.h>\n#include <stdio.h>\n#include <stdlib.h>')
 b=b.replace('\titer->off += (arc4_getword(&iter->as) % iter->skipmod) + 1;', '''\tif (iter->skipmod == 0) {
\t\tfprintf(stderr, "AUDIT_INVALID_ITERATOR skipmod=0 offset=%d usable_bits=%d\\n", iter->off, bitmap->bits);
\t\texit(2);
\t}
\titer->off += (arc4_getword(&iter->as) % iter->skipmod) + 1;''')
 assert b!=a;p.write_text(b);edits['src/iterator.c']=(a,b)
 p=build/'src/outguess.c';a=p.read_text();b=a.replace('''\tfor (where = 0; where < bits; where++) {
\t\ttmp |=''', '''\tfor (where = 0; where < bits; where++) {
\t\tif (i >= (u_int32_t)bitmap->bits) {
\t\t\tfprintf(stderr, "AUDIT_INVALID_READ offset=%u usable_bits=%d\\n", i, bitmap->bits);
\t\t\texit(2);
\t\t}
\t\ttmp |=''')
 b=b.replace('\tbuf = decode_data (tmpbuf, &datalen, as, flags);', '''\tfprintf(stderr, "AUDIT_RAW_HEADER ");
\tfor (i = 0; i < datalen; i++) fprintf(stderr, "%02x", tmpbuf[i]);
\tfprintf(stderr, "\\n");
\tbuf = decode_data (tmpbuf, &datalen, as, flags);''')
 b=b.replace('''\t\titerator_adapt(iter, bitmap, datalen);
\t\tbuf[n++]''','''\t\titerator_adapt(iter, bitmap, datalen);
\t\tif (iter->skipmod == 0)
\t\t\tfprintf(stderr, "AUDIT_REMAINING bytes=%u completed=%u\\n", datalen, n);
\t\tbuf[n++]''')
 assert 'AUDIT_INVALID_READ' in b and 'AUDIT_RAW_HEADER' in b;p.write_text(b);edits['src/outguess.c']=(a,b)
 diff=''.join(''.join(difflib.unified_diff(a.splitlines(True),b.splitlines(True),fromfile='a/'+name,tofile='b/'+name)) for name,(a,b) in edits.items())
 (out/'outguess_iterator_guard.patch').write_text(diff)
 patched=out/'outguess-guarded';compile_to(patched)
 payload=(old/'control_payload.bin').read_bytes();controls=[]
 for label,key in [('default',None),('keyed','session37-control')]:
  cover=old/f'control_outguess_{label}.jpg'
  for name,exe in [('original',baseline),('guarded',patched)]:
   dst=out/f'{name}_{label}.bin';r=extract(exe,cover,dst,key);r['exact']=dst.read_bytes()==payload;assert r['returncode']==0 and r['exact'];controls.append({'build':name,'mode':label,**r})
  # Also ensure guarded embedding remains readable by the unmodified build.
  steg=out/f'new_{label}.jpg';opts=[] if key is None else ['-k',key]
  emb=run([patched,*opts,'-d',old/'control_payload.bin',old/'control_cover.jpg',steg]);assert emb['returncode']==0,emb
  dst=out/f'cross_{label}.bin';r=extract(baseline,steg,dst,key);r['exact']=dst.read_bytes()==payload;assert r['returncode']==0 and r['exact'];controls.append({'build':'guarded embed / original extract','mode':label,**r})
 print('Six constructed recovery checks passed.',flush=True)
 prev=json.loads((old/'jpeg_file_results.json').read_text());expected={x['file']:x for x in prev['files']};rows=[]
 for p in sorted(images.glob('[0-9][0-9].jpg')):
  assert sha(p.read_bytes())==expected[p.name]['sha256'];before=extract(baseline,p,out/f'baseline_{p.stem}.bin');after=extract(patched,p,out/f'guarded_{p.stem}.bin')
  if before['returncode']==0:assert after['returncode']==0 and before['sha256']==after['sha256']
  previous=expected[p.name]['outguess'];oldhash=previous['output'].get('sha256')
  if previous['returncode']==0:assert before['returncode']==0 and before['sha256']==oldhash
  assert after['returncode']>=0 and 'runtime error:' not in after['stderr']
  rows.append({'file':p.name,'source_sha256':sha(p.read_bytes()),'previous_status':previous['returncode'],'baseline':before,'guarded':after});print(p.name,before['returncode'],'->',after['returncode'],flush=True)
 # Known no-message JPEG: white upper quarter; seeded noise below. No embedding call.
 rng=random.Random(3301038);cover=Image.new('RGB',(1024,1024),'white');noise=Image.frombytes('RGB',(1024,768),rng.randbytes(1024*768*3));cover.paste(noise,(0,256));negative=out/'unembedded_white_margin.jpg';cover.save(negative,quality=75)
 neg={'source_sha256':sha(negative.read_bytes()),'baseline':extract(baseline,negative,out/'negative_original.bin'),'guarded':extract(patched,negative,out/'negative_guarded.bin')}
 results={'source':{'original_iterator_sha256':sha(edits['src/iterator.c'][0].encode()),'original_outguess_sha256':sha(edits['src/outguess.c'][0].encode()),'patch_sha256':sha(diff.encode())},'build':{'compiler':run(['gcc','--version'])['stdout'].splitlines()[0],'flags':flags,'original_binary_sha256':sha(baseline.read_bytes()),'guarded_binary_sha256':sha(patched.read_bytes())},'diagnosis':diagnosis,'controls':controls,'files':rows,'unembedded_control':neg,'scope':'Default key only; no error correction; 75 archived files. Guards abort invalid states, do not modify traversal or encryption on successful paths.'}
 (out/'outguess_crash_summary.json').write_text(json.dumps(summarize(results),indent=2))
 (out/'outguess_crash_results.json').write_text(json.dumps(results,indent=2));print('Guarded outcomes:',collections.Counter(r['guarded']['returncode'] for r in rows));print('Unembedded control:',neg)
if __name__=='__main__':main()