import streamlit as st
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.backends import default_backend

# --- Cryptographic helper functions ---
def generate_rsa_keypair():
    """Generate a 2048-bit RSA key pair."""
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    return private_key, private_key.public_key()


def encrypt_message(pubkey, message: bytes) -> bytes:
    """Encrypt with RSA OAEP."""
    return pubkey.encrypt(
        message,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )


def decrypt_message(privkey, ciphertext: bytes) -> bytes:
    """Decrypt with RSA OAEP."""
    return privkey.decrypt(
        ciphertext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )


def derive_shared_secret(pubkey1, pubkey2) -> bytes:
    """
    Simple “shared secret”: SHA-256( server_pubkey_bytes || client_pubkey_bytes ).
    (Real SSH uses Diffie–Hellman or ECDH.)
    """
    pk1_bytes = pubkey1.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    pk2_bytes = pubkey2.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
    digest.update(pk1_bytes + pk2_bytes)
    return digest.finalize()

# --- Streamlit interface ---
def main():
    st.title("SSH‑Style Key Exchange Simulation")
    st.write("This demo simulates a basic SSH-like handshake using RSA key pairs.")

    client_msg = st.text_input("Client message:", value="Hello from client!")
    server_msg = st.text_input("Server message:", value="Hello from server!")

    if st.button("Run handshake"):
        # Generate keys
        server_priv, server_pub = generate_rsa_keypair()
        client_priv, client_pub = generate_rsa_keypair()

        # Client → Server
        enc_to_server = encrypt_message(server_pub, client_msg.encode())
        dec_at_server = decrypt_message(server_priv, enc_to_server)

        # Server → Client
        enc_to_client = encrypt_message(client_pub, server_msg.encode())
        dec_at_client = decrypt_message(client_priv, enc_to_client)

        # Shared secret derivation
        shared_secret = derive_shared_secret(server_pub, client_pub)

        st.subheader("Client → Server")
        st.write("Encrypted (first 30 bytes):", enc_to_server[:30], "...")
        st.write("Decrypted message:", dec_at_server.decode())

        st.subheader("Server → Client")
        st.write("Encrypted (first 30 bytes):", enc_to_client[:30], "...")
        st.write("Decrypted message:", dec_at_client.decode())

        st.subheader("Derived Shared Secret (SHA-256 hex)")
        st.code(shared_secret.hex())

if __name__ == "__main__":
    main()