import sys
from PIL import Image

def decode_lsb(image_path, output_path):
    print(f"[*] Extracting payload from {image_path}...")
    
    img = Image.open(image_path)
    img = img.convert('RGB')
    pixels = img.load()
    width, height = img.size
    
    # 1. Read the first 64 bits (8 bytes) to get the payload length
    header_bits = ""
    bit_idx = 0
    header_length = 64
    
    for y in range(height):
        for x in range(width):
            r, g, b = pixels[x, y]
            
            if bit_idx < header_length:
                header_bits += str(r & 1)
                bit_idx += 1
            if bit_idx < header_length:
                header_bits += str(g & 1)
                bit_idx += 1
            if bit_idx < header_length:
                header_bits += str(b & 1)
                bit_idx += 1
                
            if bit_idx >= header_length:
                break
        if bit_idx >= header_length:
            break
            
    payload_length_bytes = int(header_bits, 2)
    payload_length_bits = payload_length_bytes * 8
    print(f"[*] Detected embedded payload of size: {payload_length_bytes} bytes.")
    
    # 2. Read the actual payload
    payload_bits = ""
    bit_idx = 0
    total_bits_to_read = header_length + payload_length_bits
    
    for y in range(height):
        for x in range(width):
            r, g, b = pixels[x, y]
            
            if bit_idx >= header_length and bit_idx < total_bits_to_read:
                payload_bits += str(r & 1)
            bit_idx += 1
            
            if bit_idx >= header_length and bit_idx < total_bits_to_read:
                payload_bits += str(g & 1)
            bit_idx += 1
            
            if bit_idx >= header_length and bit_idx < total_bits_to_read:
                payload_bits += str(b & 1)
            bit_idx += 1
            
            if bit_idx >= total_bits_to_read:
                break
        if bit_idx >= total_bits_to_read:
            break
            
    # Convert bits to bytes
    payload_bytes = bytearray()
    for i in range(0, len(payload_bits), 8):
        byte = payload_bits[i:i+8]
        payload_bytes.append(int(byte, 2))
        
    # 3. Save Payload
    with open(output_path, 'wb') as f:
        f.write(payload_bytes)
        
    print(f"[+] Successfully extracted payload to {output_path}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python3 stego_decoder.py <encoded_carrier.png> <output_payload.tar.gz>")
        sys.exit(1)
        
    decode_lsb(sys.argv[1], sys.argv[2])