import os
from PIL import Image
import base64

def encrypt(plain_text, image_name):
    # Convert the plain text to bytes
    plain_bytes = plain_text.encode('utf-8')
    # Encode the bytes as base64
    b64_bytes = base64.b64encode(plain_bytes)
    # Convert the base64 bytes to a string and add padding if needed
    b64_string = b64_bytes.decode('utf-8')
    if len(b64_string) % 3 == 1:
        b64_string += '=='
    elif len(b64_string) % 3 == 2:
        b64_string += '='
    # Create an image from the base64 string
    img = Image.frombytes('RGB', (len(b64_string) // 3, 1), b64_string.encode('utf-8'))
    # Save the image as a PNG file with the specified name
    img.save(os.path.join(os.path.expanduser('~'), 'Pictures', f"{image_name}.png"))

def decrypt(image_name):
    # Open the image file
    img = Image.open(os.path.join(os.path.expanduser('~'), 'Pictures', f"{image_name}.png"))
    # Get the pixel data from the image
    pixels = img.load()
    # Create an empty string
    b64_string = ''
    # Loop through the pixels and add their RGB values to the string
    for x in range(img.width):
        r, g, b = pixels[x, 0]
        b64_string += chr(r) + chr(g) + chr(b)
    # Decode the base64 string
    b64_bytes = b64_string.encode('utf-8')
    plain_bytes = base64.b64decode(b64_bytes)
    # Convert the bytes to plain text
    plain_text = plain_bytes.decode('utf-8')
    return plain_text

def main():
    while True:
        print("Menu:")
        print("1. Encrypt plain text into an image")
        print("2. Decrypt an image into plain text")
        print("0. Quit")

        choice = input("Enter your choice: ")

        if choice == '1':
            plain_text = ""
            print("Enter the plain text (press Enter on a blank line to finish):")
            while True:
                line = input()
                if line:
                    plain_text += line + "\n"
                else:
                    break
            image_name = input("Enter a name for the image: ")
            encrypt(plain_text.strip(), image_name)
            print(f"Image '{image_name}.png' created successfully!")
        elif choice == '2':
            image_name = input("Enter the name of the image to decrypt: ")
            plain_text = decrypt(image_name)
            print("Plain text: ", plain_text)
        elif choice == '0':
            print("Exiting...")
            break
        else:
            print("Invalid choice. Please try again.\n")

if __name__ == '__main__':
    main()