#!/usr/bin/env python3
"""Cyberdelia Lazarus Recovery Workbench 0.1.0-alpha.
Read-first imaging, disk inspection, hashing and file carving.
"""
import argparse, hashlib, json, os, platform, re, struct, subprocess, sys, time
from pathlib import Path

VERSION='0.1.0-alpha'; MiB=1024**2; GiB=1024**3

def sizearg(s):
    m=re.fullmatch(r'\s*(\d+(?:\.\d+)?)\s*([kmgt]?i?b?)?\s*',s,re.I)
    if not m: raise argparse.ArgumentTypeError(s)
    n=float(m.group(1)); u=(m.group(2) or '').lower()
    mult={'':1,'b':1,'k':1024,'kb':1000,'kib':1024,'m':1024**2,'mb':1000**2,'mib':1024**2,
          'g':1024**3,'gb':1000**3,'gib':1024**3,'t':1024**4,'tb':1000**4,'tib':1024**4}
    if u not in mult: raise argparse.ArgumentTypeError(s)
    return int(n*mult[u])

def human(n):
    x=float(n)
    for u in ('B','KiB','MiB','GiB','TiB'):
        if x<1024 or u=='TiB': return f'{int(x)} B' if u=='B' else f'{x:.2f} {u}'
        x/=1024

def open_size(path,f=None,explicit=None):
    if explicit is not None: return explicit
    try:
        if os.path.isfile(path): return os.path.getsize(path)
    except OSError: pass
    close=False
    if f is None: f=open(path,'rb',buffering=0); close=True
    try:
        if sys.platform.startswith('linux'):
            try:
                import fcntl
                b=bytearray(8); fcntl.ioctl(f.fileno(),0x80081272,b,True); return struct.unpack('Q',b)[0]
            except Exception: pass
        if os.name=='nt':
            try:
                import ctypes,msvcrt
                out=ctypes.c_longlong(); ret=ctypes.c_ulong()
                ok=ctypes.windll.kernel32.DeviceIoControl(msvcrt.get_osfhandle(f.fileno()),0x0007405C,None,0,
                    ctypes.byref(out),ctypes.sizeof(out),ctypes.byref(ret),None)
                if ok: return int(out.value)
            except Exception: pass
        cur=f.tell(); f.seek(0,2); n=f.tell(); f.seek(cur); return n
    finally:
        if close: f.close()

def readat(f,off,n):
    f.seek(off); b=f.read(n)
    if len(b)!=n: raise OSError(f'short read at {off}: {len(b)}/{n}')
    return b

def fshint(f,off):
    try: b=readat(f,off,4096)
    except OSError: return 'unreadable'
    if b[3:11]==b'NTFS    ': return 'NTFS'
    if b[3:11]==b'EXFAT   ': return 'exFAT'
    if b[82:90]==b'FAT32   ': return 'FAT32'
    if b[:8]==b'-FVE-FS-': return 'BitLocker'
    return 'unknown'

def partitions(f):
    try: m=readat(f,0,512)
    except OSError: return []
    if m[510:512]!=b'\x55\xaa': return []
    out=[]; protective=False
    for i in range(4):
        e=m[446+i*16:462+i*16]; t=e[4]; start,count=struct.unpack_from('<II',e,8)
        if t and count:
            protective|=t==0xEE; out.append({'index':i+1,'start_lba':start,'sectors':count,'type':f'MBR 0x{t:02X}'})
    if protective:
        try:
            h=readat(f,512,512)
            if h[:8]==b'EFI PART':
                lba=struct.unpack_from('<Q',h,72)[0]; count=min(struct.unpack_from('<I',h,80)[0],4096); es=struct.unpack_from('<I',h,84)[0]
                raw=readat(f,lba*512,count*es); out=[]
                for i in range(count):
                    e=raw[i*es:(i+1)*es]
                    if e[:16]==b'\0'*16: continue
                    a=struct.unpack_from('<Q',e,32)[0]; z=struct.unpack_from('<Q',e,40)[0]
                    if z<a: continue
                    name=e[56:128].decode('utf-16le','ignore').rstrip('\0')
                    out.append({'index':len(out)+1,'start_lba':a,'sectors':z-a+1,'type':'GPT','name':name})
        except OSError: pass
    for p in out: p['offset']=p['start_lba']*512; p['fs']=fshint(f,p['offset'])
    return out

def cmd_devices(a):
    if platform.system()=='Windows':
        q='Get-CimInstance Win32_DiskDrive | Select Index,Model,SerialNumber,Size,InterfaceType,DeviceID | ConvertTo-Json -Compress'
        try: print(json.dumps(json.loads(subprocess.check_output(['powershell','-NoProfile','-Command',q],text=True)),indent=2))
        except Exception as e: print(json.dumps({'error':str(e)},indent=2))
    elif platform.system()=='Linux':
        try: print(subprocess.check_output(['lsblk','-J','-b','-o','NAME,PATH,SIZE,MODEL,SERIAL,TYPE,FSTYPE,MOUNTPOINTS'],text=True))
        except Exception as e: print(json.dumps({'error':str(e)},indent=2))
    else: print(json.dumps({'error':'device listing currently supports Windows and Linux'}))

def cmd_inspect(a):
    with open(a.source,'rb',buffering=0) as f:
        n=open_size(a.source,f,a.size); print(f'Source: {a.source}\nSize: {n} ({human(n)})')
        root=fshint(f,0)
        if root!='unknown': print('Filesystem at byte 0:',root)
        ps=partitions(f)
        if not ps: print('No MBR/GPT partitions detected.'); return
        for p in ps: print(f"[{p['index']}] offset={p['offset']} ({human(p['offset'])}) size={human(p['sectors']*512)} fs={p['fs']} {p.get('name') or p['type']}")

def journal(j,obj):
    j.write(json.dumps({'ts':time.time(),**obj},sort_keys=True)+'\n'); j.flush(); os.fsync(j.fileno())

def recover(src,dst,off,n,minblock,p,j,bad):
    try:
        src.seek(off); b=src.read(n)
        if len(b)!=n: raise OSError(f'short read {len(b)}/{n}')
        dst.seek(off); dst.write(b); journal(j,{'pass':p,'status':'good','offset':off,'length':n}); return
    except OSError as e:
        if n<=minblock:
            bad.append((off,n)); journal(j,{'pass':p,'status':'bad','offset':off,'length':n,'error':str(e)}); return
    half=max(minblock,(n//2//minblock)*minblock)
    if half<=0 or half>=n: half=n//2
    recover(src,dst,off,half,minblock,p,j,bad); recover(src,dst,off+half,n-half,minblock,p,j,bad)

def cmd_clone(a):
    if os.path.abspath(a.source)==os.path.abspath(a.destination): raise SystemExit('source and destination must differ')
    if a.destination.lower().startswith(r'\\.\physicaldrive'): raise SystemExit('v0.1 refuses raw-device destinations; use an image file')
    dest=Path(a.destination); mp=Path(a.map or (a.destination+'.lazarus-map.jsonl'))
    if dest.exists() and not a.resume: raise SystemExit('destination exists; use --resume or another path')
    cursor=0
    if a.resume and mp.exists():
        for line in mp.read_text(errors='ignore').splitlines():
            try:
                o=json.loads(line)
                if o.get('status')=='cursor': cursor=max(cursor,int(o.get('offset',0)))
            except Exception: pass
    with open(a.source,'rb',buffering=0) as src, open(dest,'r+b' if dest.exists() else 'w+b',buffering=0) as dst, mp.open('a') as j:
        total=open_size(a.source,src,a.size)
        if not dest.stat().st_size: dst.truncate(total)
        if dest.stat().st_size!=total: raise SystemExit('destination size mismatch')
        bad=[]; pos=cursor; journal(j,{'event':'start','source':a.source,'destination':a.destination,'size':total,'cursor':cursor})
        while pos<total:
            n=min(a.chunk,total-pos)
            try:
                src.seek(pos); b=src.read(n)
                if len(b)!=n: raise OSError(f'short read {len(b)}/{n}')
                dst.seek(pos); dst.write(b); journal(j,{'pass':0,'status':'good','offset':pos,'length':n})
            except OSError as e:
                bad.append((pos,n)); journal(j,{'pass':0,'status':'bad','offset':pos,'length':n,'error':str(e)})
            pos+=n; journal(j,{'status':'cursor','offset':pos})
            print(f'\rPass 0 {pos/total:6.2%} {human(pos)}/{human(total)} bad={human(sum(n for _,n in bad))}',end='',flush=True)
        print(); dst.flush(); os.fsync(dst.fileno())
        cur=bad
        for p in range(1,a.retry_passes+1):
            if not cur: break
            nxt=[]; print(f'Retry pass {p}: {len(cur)} ranges')
            for i,(off,n) in enumerate(cur,1):
                recover(src,dst,off,n,a.sector,p,j,nxt); print(f'\r{i}/{len(cur)} unresolved={human(sum(x[1] for x in nxt))}',end='',flush=True)
            print(); cur=nxt; dst.flush(); os.fsync(dst.fileno())
        summary={'source':a.source,'destination':a.destination,'size':total,'unrecovered_bytes':sum(n for _,n in cur),
                 'unrecovered_ranges':[{'offset':o,'length':n} for o,n in cur]}
        Path(str(mp)+'.summary.json').write_text(json.dumps(summary,indent=2)+'\n'); journal(j,{'event':'complete','unrecovered_bytes':summary['unrecovered_bytes']})
        print('Clone complete. Unrecovered:',human(summary['unrecovered_bytes']))

def until(f,off,marker,maxsize):
    f.seek(off); out=bytearray()
    while len(out)<maxsize:
        b=f.read(min(4*MiB,maxsize-len(out)))
        if not b: break
        old=max(0,len(out)-len(marker)-8); out.extend(b); i=out.find(marker,old)
        if i>=0: return bytes(out[:i+len(marker)])

def zip_candidate(f,off,maxsize):
    f.seek(off); out=bytearray(); mark=b'PK\x05\x06'
    while len(out)<maxsize:
        b=f.read(min(4*MiB,maxsize-len(out)))
        if not b: break
        old=max(0,len(out)-32); out.extend(b); i=out.find(mark,old)
        if i>=0 and len(out)>=i+22:
            com=struct.unpack_from('<H',out,i+20)[0]; need=i+22+com
            while len(out)<need:
                b=f.read(min(4*MiB,need-len(out)))
                if not b: break
                out.extend(b)
            if len(out)>=need: return bytes(out[:need])

def classify_zip(b):
    if b'word/document.xml' in b: return '.docx'
    if b'xl/workbook.xml' in b: return '.xlsx'
    if b'ppt/presentation.xml' in b: return '.pptx'
    return '.zip'

def cmd_carve(a):
    specs={'jpg':(b'\xff\xd8\xff',b'\xff\xd9','.jpg',128*MiB),'png':(b'\x89PNG\r\n\x1a\n',b'\0\0\0\0IEND\xaeB`\x82','.png',128*MiB),
           'pdf':(b'%PDF-',b'%%EOF','.pdf',512*MiB),'zip':(b'PK\x03\x04',None,'.zip',2*GiB)}
    wanted=[x.strip().lower() for x in a.types.split(',') if x.strip()]
    if any(x not in specs for x in wanted): raise SystemExit('types: jpg,png,pdf,zip')
    out=Path(a.output); out.mkdir(parents=True,exist_ok=True); mf=(out/'carve-manifest.jsonl').open('a')
    with open(a.source,'rb',buffering=0) as f:
        total=open_size(a.source,f,a.size); tail=b''; pos=0; seen=set(); count=0; maxsig=max(len(specs[t][0]) for t in wanted)
        while pos<total:
            f.seek(pos); data=f.read(min(a.block,total-pos))
            if not data: break
            buf=tail+data; base=pos-len(tail)
            hits=[]
            for t in wanted:
                sig=specs[t][0]; s=0
                while True:
                    i=buf.find(sig,s)
                    if i<0: break
                    off=base+i
                    if (off,t) not in seen: seen.add((off,t)); hits.append((off,t))
                    s=i+1
            for off,t in sorted(hits):
                start,end,ext,limit=specs[t]; limit=min(limit,a.max_file)
                try: b=zip_candidate(f,off,limit) if t=='zip' else until(f,off,end,limit)
                except OSError: b=None
                if not b: continue
                if t=='pdf': b=b.rstrip(b'\0')+b'\n'
                if t=='zip': ext=classify_zip(b)
                h=hashlib.sha256(b).hexdigest(); name=f'{off:016x}_{h[:12]}{ext}'; (out/name).write_bytes(b)
                mf.write(json.dumps({'offset':off,'type':t,'bytes':len(b),'sha256':h,'path':name})+'\n'); mf.flush(); count+=1
                print(f'[{count}] {human(off):>10} {t:<4} {human(len(b)):>10} {name}')
            tail=buf[-(maxsig-1):]; pos+=len(data)
    mf.close(); print(f'Carved {count} files. Manifest: {out/"carve-manifest.jsonl"}')

def cmd_hash(a):
    h=hashlib.sha256(); total=open_size(a.source); done=0
    with open(a.source,'rb',buffering=0) as f:
        while True:
            b=f.read(a.chunk)
            if not b: break
            h.update(b); done+=len(b); print(f'\r{done/total:6.2%} {human(done)}/{human(total)}',end='',flush=True)
    print('\n'+h.hexdigest())

def parser():
    p=argparse.ArgumentParser(prog='lazarus',description='Cyberdelia Lazarus: read-first disk imaging, inspection and recovery')
    p.add_argument('--version',action='version',version=f'%(prog)s {VERSION}'); s=p.add_subparsers(dest='cmd',required=True)
    x=s.add_parser('devices'); x.set_defaults(fn=cmd_devices)
    x=s.add_parser('inspect'); x.add_argument('source'); x.add_argument('--size',type=sizearg); x.set_defaults(fn=cmd_inspect)
    x=s.add_parser('clone'); x.add_argument('source'); x.add_argument('destination'); x.add_argument('--map'); x.add_argument('--size',type=sizearg); x.add_argument('--chunk',type=sizearg,default=4*MiB); x.add_argument('--sector',type=sizearg,default=512); x.add_argument('--retry-passes',type=int,default=1); x.add_argument('--resume',action='store_true'); x.set_defaults(fn=cmd_clone)
    x=s.add_parser('carve'); x.add_argument('source'); x.add_argument('output'); x.add_argument('--types',default='jpg,png,pdf,zip'); x.add_argument('--size',type=sizearg); x.add_argument('--block',type=sizearg,default=8*MiB); x.add_argument('--max-file',type=sizearg,default=2*GiB); x.set_defaults(fn=cmd_carve)
    x=s.add_parser('hash'); x.add_argument('source'); x.add_argument('--chunk',type=sizearg,default=8*MiB); x.set_defaults(fn=cmd_hash)
    return p

def main():
    p=parser(); a=p.parse_args()
    try: a.fn(a); return 0
    except PermissionError: print('Permission denied. Raw disks normally require Administrator/root.',file=sys.stderr); return 13
    except KeyboardInterrupt: print('\nInterrupted. Clone map/output remain for resume.',file=sys.stderr); return 130
if __name__=='__main__': raise SystemExit(main())