import argparse
import os
import struct
from Crypto.Cipher import AES
from PIL import Image

def encrypt_file(key, in_filename, out_filename=None, chunksize=64*1024):
    if not out_filename:
        out_filename = in_filename + '.encrypted'

    iv = os.urandom(16)
    key_length = len(key)
    if key_length not in (16, 24, 32):
        if key_length < 16:
            key = key.ljust(16, '\0')
        elif key_length < 24:
            key = key.ljust(24, '\0')
        else:
            key = key.ljust(32, '\0')
    encryptor = AES.new(key, AES.MODE_CBC, iv)

    filesize = os.path.getsize(in_filename)

    with open(in_filename, 'rb') as infile:
        with open(out_filename, 'wb') as outfile:
            outfile.write(struct.pack('<Q', filesize))
            outfile.write(iv)

            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    chunk += b' ' * (16 - len(chunk) % 16)

                outfile.write(encryptor.encrypt(chunk))

def decrypt_file(key, in_filename, out_filename=None, chunksize=24*1024):
    if not out_filename:
        out_filename = os.path.splitext(in_filename)[0]

    with open(in_filename, 'rb') as infile:
        filesize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
        iv = infile.read(16)
        key_length = len(key)
        if key_length not in (16, 24, 32):
            if key_length < 16:
                key = key.ljust(16, '\0')
            elif key_length < 24:
                key = key.ljust(24, '\0')
            else:
                key = key.ljust(32, '\0')
        decryptor = AES.new(key, AES.MODE_CBC, iv)

        with open(out_filename, 'wb') as outfile:
            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                outfile.write(decryptor.decrypt(chunk))

            outfile.truncate(filesize)

def hide_text_in_image(text, image_filename, output_filename=None):
    image = Image.open(image_filename)
    pixels = image.load()
    width, height = image.size

    text_length = len(text)
    text_length = int(text_length)
    pixels[0, 0] = text_length

    char_index = 0
    for y in range(height):
        for x in range(width):
            if char_index < text_length:
                char = text[char_index]
                ascii_code = ord(char)
                pixel = pixels[x, y]
                if isinstance(pixel, tuple) and len(pixel) == 3:
                    red, green, blue = pixel
                else:
                    continue
                red = (red & 0xFE) | ((ascii_code >> 7) & 0x01)
                green = (green & 0xFE) | ((ascii_code >> 6) & 0x01)
                blue = (blue & 0xFE) | ((ascii_code >> 5) & 0x01)
                pixels[x, y] = (red, green, blue)
                char_index += 1
            else:
                break

    if not output_filename:
        output_filename = os.path.splitext(image_filename)[0] + '_hidden.png'
    image.save(output_filename)


def extract_key_from_image(image_filename):
    image = Image.open(image_filename)
    pixels = image.load()
    width, height = image.size

    text_length = pixels[0, 0][0]

    text = ''
    char_index = 0
    for y in range(height):
        for x in range(width):
            if char_index < text_length:
                pixel = pixels[x, y]
                if isinstance(pixel, tuple) and len(pixel) == 3:
                    red, green, blue = pixel
                else:
                    continue
                ascii_code = ((red & 0x01) << 7) | ((green & 0x01) << 6) | ((blue & 0x01) << 5)
                char = chr(ascii_code)
                text += char
                char_index += 1
            else:
                break

    return text

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Скрипт для скрытия и извлечения текста в изображениях методом стенографии.")
    parser.add_argument("command", choices=["encrypt_file", "decrypt_file", "encrypt_image", "decrypt_image"])
    parser.add_argument("input", help="Входной файл")
    parser.add_argument("--output", help="Выходной файл (если не указан, будет использоваться стандартное имя)")
    args = parser.parse_args()

    key = extract_key_from_image(args.input)

    if args.command == "encrypt_file":
        encrypt_file(key, args.input, args.output)
    elif args.command == "decrypt_file":
        decrypt_file(key, args.input, args.output)
    elif args.command == "encrypt_image":
        text = input("Введите текст, который нужно скрыть в изображении: ")
        hide_text_in_image(text, args.input, args.output)
    elif args.command == "decrypt_image":
        print(extract_key_from_image(args.input))