import sys, hmac, struct, hashlib
from PIL import Image

def text_to_bits(text):
    bits = bin(int.from_bytes(text.encode('utf-8'), 'big'))[2:]
    return bits.zfill(8 * ((len(bits) + 7) // 8))

def bits_to_text(bits):
    n = int(bits, 2)
    return n.to_bytes((n.bit_length() + 7) // 8, 'big').decode('utf-8', 'ignore')

def get_pixel_path(image_size, password):
    width, height = image_size
    coords = list(range(width * height))
    key = password.encode('utf-8')
    
    def get_random_bytes(counter):
        return hmac.new(key, struct.pack('>Q', counter), digestmod=hashlib.sha256).digest()
    
    # Kryptografischer Fisher-Yates Shuffle mit Rejection Sampling
    counter = 0
    byte_buffer = b''
    for i in range(len(coords) - 1, 0, -1):
        upper = i + 1
        limit = (0x100000000 // upper) * upper
        
        while True:
            while len(byte_buffer) < 4:
                byte_buffer += get_random_bytes(counter)
                counter += 1
            val = struct.unpack('>I', byte_buffer[:4])[0]
            byte_buffer = byte_buffer[4:]
            
            if val < limit:
                j = val % upper
                break
                
        coords[i], coords[j] = coords[j], coords[i]
    
    return [(c % width, c // width) for c in coords]

def embed_data(image_path, output_path, payload, password):
    img = Image.open(image_path).convert('RGB')
    pixels = img.load()
    binary_msg = text_to_bits(payload + "###EOF###")
    
    # Padding, falls Bit-Länge nicht exakt durch 3 teilbar ist
    remainder = len(binary_msg) % 3
    if remainder != 0:
        binary_msg += '0' * (3 - remainder)
        
    coords = get_pixel_path(img.size, password)
    
    if (len(binary_msg) // 3) > len(coords): 
        raise ValueError("Bild zu klein für diese Datenmenge!")
        
    for i in range(len(binary_msg) // 3):
        x, y = coords[i]
        r, g, b = pixels[x, y]
        bit_pos = i * 3
        
        pixels[x, y] = (
            (r & ~1) | int(binary_msg[bit_pos]),
            (g & ~1) | int(binary_msg[bit_pos + 1]),
            (b & ~1) | int(binary_msg[bit_pos + 2])
        )
        
    img.save(output_path, "PNG")
    print(f"Daten erfolgreich eingebettet und gespeichert unter: {output_path}")

def extract_data(image_path, password):
    img = Image.open(image_path).convert('RGB')
    pixels = img.load()
    coords = get_pixel_path(img.size, password)
    
    extracted_bits = ""
    extracted_text = ""
    
    for x, y in coords:
        r, g, b = pixels[x, y]
        extracted_bits += str(r & 1) + str(g & 1) + str(b & 1)
        
        while len(extracted_bits) >= 8:
            byte_str = extracted_bits[:8]
            extracted_bits = extracted_bits[8:]
            extracted_text += bits_to_text(byte_str)
            
            if extracted_text.endswith("###EOF###"): 
                return extracted_text[:-9]
                
    return "Fehler: Daten nicht gefunden oder falsches Passwort."

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Nutzung: python3 stegano_otp.py [embed|extract] ...")
    else:
        mode = sys.argv[1]
        if mode == "embed" and len(sys.argv) == 6:
            embed_data(sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5])
        elif mode == "extract" and len(sys.argv) == 4:
            print(extract_data(sys.argv[2], sys.argv[3]))