from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.backends import default_backend
import os
import base64
import getpass

def init_keys(private_key_file="private_key.pem", public_key_file="public_key.pem"):
    """
    Generate a new RSA private/public key pair and save them to files.
    """
    # Check if keys already exist
    if os.path.exists(private_key_file) and os.path.exists(public_key_file):
        print("Keys already exist. Skipping key generation.")
        return

    # Generate a private key
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )

    # Generate the public key
    public_key = private_key.public_key()

    # Serialize and save the private key
    with open(private_key_file, "wb") as f:
        f.write(
            private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
        )

    # Serialize and save the public key
    with open(public_key_file, "wb") as f:
        f.write(
            public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )
        )

    print(f"Private key saved to {private_key_file}")
    print(f"Public key saved to {public_key_file}")


def encrypt_passphrase(passphrase, public_key_file="public_key.pem"):
    """
    Encrypt a passphrase using the public key.
    Returns the encrypted passphrase as a base64-encoded string.
    """
    # Load the public key
    with open(public_key_file, "rb") as f:
        public_key = serialization.load_pem_public_key(
            f.read(),
            backend=default_backend()
        )

    # Encrypt the passphrase
    encrypted = public_key.encrypt(
        passphrase.encode(),
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

    # Return the encrypted passphrase as a base64-encoded string
    return base64.b64encode(encrypted).decode()

def decrypt_passphrase(encrypted_passphrase, private_key_file="private_key.pem"):
    """
    Decrypt an encrypted passphrase using the private key.
    Returns the decrypted passphrase as a string.
    """
    # Load the private key
    with open(private_key_file, "rb") as f:
        private_key = serialization.load_pem_private_key(
            f.read(),
            password=None,
            backend=default_backend()
        )

    # Decode the base64-encoded encrypted passphrase
    encrypted = base64.b64decode(encrypted_passphrase.encode())

    # Decrypt the passphrase
    decrypted = private_key.decrypt(
        encrypted,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

    # Return the decrypted passphrase
    return decrypted.decode()

def create_luks_mount_config(config_file="luks_mount.conf", public_key_file="public_key.pem"):
    """
    Create a new luks_mount.conf file by interactively prompting for device, mount point, and password.
    If the configuration file already exists, provide an option to quit or append to the existing file.
    """
    mode = "w"
    if os.path.exists(config_file):
        print(f"Configuration file '{config_file}' already exists.")
        choice = input("Do you want to (a)ppend to it or (q)uit? ").strip().lower()
        if choice == 'q':
            print("Exiting without making changes.")
            return
        elif choice == 'a':
            mode = "a"
        else:
            print("Invalid choice. Exiting without making changes.")
            return

    with open(config_file, mode) as f:
        if mode == "w":
            f.write("# LUKS mount configuration file\n")
            f.write("# Format: UUID mount_point encrypted_passphrase\n")

        while True:
            device = input("Enter the device UUID (e.g., UUID): ").strip()
            mount_point = input("Enter the mount point (e.g., /media/<username>/luks-<UUID>): ").strip()
            passphrase = getpass.getpass("Enter the passphrase: ").strip()

            encrypted_passphrase = encrypt_passphrase(passphrase, public_key_file)
            f.write(f"{device} {mount_point} {encrypted_passphrase}\n")

            another = input("Do you want to add another entry? (y/n): ").strip().lower()
            if another != 'y':
                break

    print(f"Configuration saved to {config_file}")


# Example usage
if __name__ == "__main__":
    # Initialize keys (run this once to generate keys)
    init_keys()

    # Encrypt a passphrase
    #passphrases = ["5khbS4JFN25tPVrkzv2b2Q==", "kT3vPK09f+KhNWp7LpU1jg==", "3fhbN/Nt1YFPM5AUbJ4ERg=="]
    passphrases= ["df638b83d86404800bd631cf3ca3a483", "f1c873af36003777dbf7fae275e5f159", "23d828a6cbde48cd826e862da6b3e3e9"]
    
    print(decrypt_passphrase('fAqe02nfx1L+YZK/kLFbsd1nVCl5w/LJxxMeSRXLsEm+vq49dRI4eYSzjl3w2YqyPb4JeVXfMDWCTm523k9e/KI+GF8zAPisI1icWM/k+yzMCrW9Ga7rzwW092Uepm7IU4z7LwU4Z9t2wmQkeb6ulso6lcaYzVoCcIHRJ0gCkBZBb0nA1lXQvf2UNFfe9kkzX+DR1MCovY9SAsIndisCh6y0IxEBbfCTASU5VYIcngiKayS6flVtGuLOz/3S6Z+T5GDN7dLC97wf/yVBQJ2lPdgTcawv2qV59OcxDGUeHjwi340+UIramslrxmGpFMyBkmLYqcjKkWsVD6HF2msu6w=='))
    print(decrypt_passphrase('huZoCnJBUay3XSZt36XoBqckexy5iYB6LT41Wscxx/8KujJQWas/hwn6tdlr7ss8pTs2LNW71CvcRdHx6p8YgBs0aYV5o4z8P5yLahGH0ZgmA/Tc+brQN3+/s1cyFNQ8LZK0D2jizpah5Np469EO5hLT//lx4/WEobX79CS1WRZxetbqmgLqKrFcbBh50BYS/Ka61P63IIlHXCF7GySPZM1m0bRfphmQ/93FPngkyinxK0Sho0i7ddXC+bwXWeZSAi9nPTNZ4tGBa7hVzJBNVmoI4zIdjXtfJeGQ21bSPIPI+KsfE3/3Oe2AFmU1o8axbxNQJGbmee6BqJIhArzAzA=='))

    for passphrase in passphrases:
         encrypted = encrypt_passphrase(passphrase)
         print(f"Encrypted passphrase: {encrypted}")

         # Decrypt the passphrase
         decrypted = decrypt_passphrase(encrypted)
         print(f"Decrypted passphrase: {decrypted}")