import sys
import os
import math
import qrcode

def generate_qr_archive(payload_path, output_dir, chunk_size=2000):
    print(f"[*] Generating Machine-Readable QR Archive from {payload_path}...")
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    with open(payload_path, 'rb') as f:
        payload_bytes = f.read()
        
    total_bytes = len(payload_bytes)
    total_chunks = math.ceil(total_bytes / chunk_size)
    print(f"[*] Payload size: {total_bytes} bytes. Splitting into {total_chunks} chunks of {chunk_size} bytes.")
    
    for i in range(total_chunks):
        start_idx = i * chunk_size
        end_idx = min(start_idx + chunk_size, total_bytes)
        chunk_data = payload_bytes[start_idx:end_idx]
        
        # We need to encode the binary data using Base64 or ISO-8859-1 for QR code compatibility
        import base64
        chunk_b64 = base64.b64encode(chunk_data).decode('utf-8')
        
        header = f"[CHUNK_{i+1:04d}_OF_{total_chunks:04d}]\n"
        qr_data = header + chunk_b64
        
        qr = qrcode.QRCode(
            version=None,
            error_correction=qrcode.constants.ERROR_CORRECT_L,
            box_size=10,
            border=4,
        )
        qr.add_data(qr_data)
        qr.make(fit=True)
        
        img = qr.make_image(fill_color="black", back_color="white")
        
        filename = f"chunk_{i+1:04d}.png"
        img.save(os.path.join(output_dir, filename))
        
        if (i+1) % 100 == 0 or (i+1) == total_chunks:
            print(f"  -> Generated {i+1}/{total_chunks} QR codes...")
            
    print(f"[+] QR Archive successfully generated in {output_dir}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python3 generate_qr_archive.py <payload.tar.gz> <output_dir>")
        sys.exit(1)
        
    generate_qr_archive(sys.argv[1], sys.argv[2])