import sys
import os
from PIL import Image

def encode_lsb(image_path, payload_path, output_path):
    print(f"[*] Encoding {payload_path} into {image_path}...")
    
    # 1. Read Payload
    with open(payload_path, 'rb') as f:
        payload_bytes = f.read()
        
    # Inject header to denote length of the file (8 bytes / 64 bits)
    payload_length = len(payload_bytes)
    header = payload_length.to_bytes(8, byteorder='big')
    full_payload = header + payload_bytes
    
    # Convert to binary string
    binary_payload = ''.join([format(b, '08b') for b in full_payload])
    payload_len_bits = len(binary_payload)
    print(f"[*] Payload is {payload_length} bytes ({payload_len_bits} bits).")

    # 2. Open Image
    img = Image.open(image_path)
    img = img.convert('RGB')
    pixels = img.load()
    width, height = img.size
    
    max_bits = width * height * 3
    if payload_len_bits > max_bits:
        raise ValueError(f"Payload too large! Need {payload_len_bits} bits, but image can only hold {max_bits} bits.")

    # 3. Inject Data via LSB
    bit_idx = 0
    for y in range(height):
        for x in range(width):
            if bit_idx >= payload_len_bits:
                break
                
            r, g, b = pixels[x, y]
            
            if bit_idx < payload_len_bits:
                r = (r & ~1) | int(binary_payload[bit_idx])
                bit_idx += 1
            if bit_idx < payload_len_bits:
                g = (g & ~1) | int(binary_payload[bit_idx])
                bit_idx += 1
            if bit_idx < payload_len_bits:
                b = (b & ~1) | int(binary_payload[bit_idx])
                bit_idx += 1
                
            pixels[x, y] = (r, g, b)
            
        if bit_idx >= payload_len_bits:
            break

    # 4. Save
    img.save(output_path, format="PNG")
    print(f"[+] Successfully saved encoded carrier to {output_path}")

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python3 stego_encoder.py <payload.tar.gz> <carrier.png> <output.png>")
        sys.exit(1)
        
    encode_lsb(sys.argv[2], sys.argv[1], sys.argv[3])