# generate_keys.py
# Run this script first. It creates the key pairs for both the Sender and the Receiver and saves them as .pem files.
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

# 1. Generate Sender's Keys
sender_private = rsa.generate_private_key(public_exponent=65537, key_size=2048)
sender_public = sender_private.public_key()

with open("sender_private_key.pem", "wb") as f:
    f.write(
        sender_private.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )

with open("sender_public_key.pem", "wb") as f:
    f.write(
        sender_public.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
    )

# 2. Generate Receiver's Keys
receiver_private = rsa.generate_private_key(
    public_exponent=65537, key_size=2048
)
receiver_public = receiver_private.public_key()

with open("receiver_private_key.pem", "wb") as f:
    f.write(
        receiver_private.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )

with open("receiver_public_key.pem", "wb") as f:
    f.write(
        receiver_public.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
    )

print(" All RSA keys generated and saved successfully:")
print("  - sender_private_key.pem & sender_public_key.pem")
print("  - receiver_private_key.pem & receiver_public_key.pem")





# Step 2: sender.py
# Run this script second. It loads the keys, encrypts the secret message for the receiver, signs it, and saves encrypted_message.bin and signature.bin.
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

# 1. Load Receiver's Public Key (used to ENCRYPT)
with open("receiver_public_key.pem", "rb") as f:
    receiver_public_key = serialization.load_pem_public_key(f.read())

# 2. Load Sender's Private Key (used to SIGN)
with open("sender_private_key.pem", "rb") as f:
    sender_private_key = serialization.load_pem_private_key(
        f.read(), password=None
    )

# 3. Get message dynamically from user input
user_input = input("Enter the message you want to encrypt and sign: ")
message = user_input.encode("utf-8")  # Convert string to bytes

# 4. Encrypt the message with Receiver's Public Key
encrypted_message = receiver_public_key.encrypt(
    message,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=hashes.SHA256()),
        algorithm=hashes.SHA256(),
        label=None,
    ),
)

# 5. Sign the message with Sender's Private Key
signature = sender_private_key.sign(
    message,
    padding.PSS(
        mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
    ),
    hashes.SHA256(),
)

# 6. Save outputs to disk
with open("encrypted_message.bin", "wb") as f:
    f.write(encrypted_message)

with open("signature.bin", "wb") as f:
    f.write(signature)

print("\n Message encrypted and signed successfully.")
print("Saved 'encrypted_message.bin' and 'signature.bin'.")





# Step 3: receiver.py
# Run this script third. It loads the files, decrypts the message using the receiver's private key, and verifies the sender's signature.
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

# 1. Load Receiver's Private Key (used to DECRYPT)
with open("receiver_private_key.pem", "rb") as f:
    receiver_private_key = serialization.load_pem_private_key(
        f.read(), password=None
    )

# 2. Load Sender's Public Key (used to VERIFY SIGNATURE)
with open("sender_public_key.pem", "rb") as f:
    sender_public_key = serialization.load_pem_public_key(f.read())

# 3. Load Encrypted Message and Signature from files
with open("encrypted_message.bin", "rb") as f:
    encrypted_message = f.read()

with open("signature.bin", "rb") as f:
    signature = f.read()

# 4. Decrypt the message
decrypted_bytes = receiver_private_key.decrypt(
    encrypted_message,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=hashes.SHA256()),
        algorithm=hashes.SHA256(),
        label=None,
    ),
)

original_message = decrypted_bytes.decode("utf-8")
print(f"\n Decrypted Message: '{original_message}'")

# DYNAMIC TAMPERING PROMPT
print("\n--- Signature Verification Options ---")
choice = input("Do you want to tamper with the message? (y/n): ").strip().lower()

if choice == "y":
    altered_text = input("Enter the new/altered message: ")
    message_to_verify = altered_text.encode("utf-8")
else:
    message_to_verify = decrypted_bytes

# 5. Verify the signature
try:
    sender_public_key.verify(
        signature,
        message_to_verify,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH,
        ),
        hashes.SHA256(),
    )
    print("\n Signature Verified: Message is authentic and untampered!")
except InvalidSignature:
    print(
        "\n SECURITY ALERT: Signature Verification Failed! The message was tampered with!"
    )
