public.py

import base64, os
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


# =============================
# AES-256-GCM ENCRYPT
# =============================
def encrypt_message(message: str, passphrase: str) -> str:
    message_bytes = message.encode()
    passphrase_bytes = passphrase.encode()

    salt = os.urandom(16)     # 16 bytes salt (NEW every time)
    nonce = os.urandom(12)    # 12 bytes nonce (NEW every time)

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100_000,
    )
    key = kdf.derive(passphrase_bytes)

    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, message_bytes, None)

    # Combine → salt + nonce + ciphertext
    blob = salt + nonce + ciphertext
    encrypted_text = base64.urlsafe_b64encode(blob).decode()

    return encrypted_text


# =============================
# AES-256-GCM DECRYPT
# =============================
def decrypt_message(encrypted_text: str, passphrase: str) -> str:
    passphrase_bytes = passphrase.encode()

    raw = base64.urlsafe_b64decode(encrypted_text)

    salt = raw[:16]          # first 16 bytes
    nonce = raw[16:28]       # next 12 bytes
    ciphertext = raw[28:]    # rest

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100_000,
    )
    key = kdf.derive(passphrase_bytes)

    aesgcm = AESGCM(key)
    decrypted_bytes = aesgcm.decrypt(nonce, ciphertext, None)

    return decrypted_bytes.decode()


# =============================
# MENU SYSTEM
# =============================
def main():
    print("\n==============================")
    print("  AES-256-GCM Encryption Tool")
    print("==============================")
    print("1) Encrypt a Message")
    print("2) Decrypt a Message")
    print("3) Exit")
    print("==============================")

    choice = input("Choose an option: ").strip()

    if choice == "1":
        msg = input("\nEnter message to encrypt:\n> ")
        key = input("Enter a strong passphrase:\n> ")
        encrypted = encrypt_message(msg, key)

        print("\n🔒 ENCRYPTED TEXT (Share / Post / Send):")
        print(encrypted)

        print("\n🔑 IMPORTANT:")
        print("Only your passphrase can decrypt this text.")

    elif choice == "2":
        encrypted = input("\nPaste the encrypted text:\n> ")
        key = input("Enter the passphrase:\n> ")

        try:
            decrypted = decrypt_message(encrypted, key)
            print("\n🔓 DECRYPTED MESSAGE:")
            print(decrypted)
        except Exception as e:
            print("\n❌ ERROR: Wrong passphrase OR corrupted encrypted text.")

    else:
        print("Goodbye!")


if __name__ == "__main__":
    main()