Cryptography and Encryption Samples

# Symmetric Encryption with AES
# Required: pip install cryptography pycryptodome

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding, hashes
from cryptography.hazmat.backends import default_backend
import os
import base64

class AESCipher:
    """AES Encryption/Decryption utility class"""

    def __init__(self, key: bytes = None):
        """Initialize with a key or generate a new one"""
        self.key = key or os.urandom(32)  # 256-bit key for AES-256
        self.iv = None

    def generate_key(self) -> bytes:
        """Generate a new random AES key"""
        return os.urandom(32)

    def encrypt_cbc(self, plaintext: str) -> dict:
        """Encrypt using AES-CBC mode"""
        # Generate random IV
        iv = os.urandom(16)

        # Pad the plaintext
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(plaintext.encode()) + padder.finalize()

        # Encrypt
        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()

        return {
            'ciphertext': base64.b64encode(ciphertext).decode(),
            'iv': base64.b64encode(iv).decode(),
            'mode': 'CBC'
        }

    def decrypt_cbc(self, ciphertext: str, iv: str) -> str:
        """Decrypt AES-CBC encrypted data"""
        ciphertext_bytes = base64.b64decode(ciphertext)
        iv_bytes = base64.b64decode(iv)

        # Decrypt
        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv_bytes), backend=default_backend())
        decryptor = cipher.decryptor()
        padded_plaintext = decryptor.update(ciphertext_bytes) + decryptor.finalize()

        # Unpad
        unpadder = padding.PKCS7(128).unpadder()
        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()

        return plaintext.decode()

    def encrypt_gcm(self, plaintext: str, associated_data: str = None) -> dict:
        """Encrypt using AES-GCM mode (authenticated encryption)"""
        # Generate random IV and nonce
        iv = os.urandom(12)

        # Encrypt with GCM
        encryptor = Cipher(
            algorithms.AES(self.key),
            modes.GCM(iv),
            backend=default_backend()
        ).encryptor()

        # Add associated data if provided
        if associated_data:
            encryptor.authenticate_additional_data(associated_data.encode())

        ciphertext = encryptor.update(plaintext.encode()) + encryptor.finalize()

        return {
            'ciphertext': base64.b64encode(ciphertext).decode(),
            'iv': base64.b64encode(iv).decode(),
            'tag': base64.b64encode(encryptor.tag).decode(),
            'mode': 'GCM',
            'associated_data': associated_data
        }

    def decrypt_gcm(self, ciphertext: str, iv: str, tag: str, associated_data: str = None) -> str:
        """Decrypt AES-GCM encrypted data"""
        ciphertext_bytes = base64.b64decode(ciphertext)
        iv_bytes = base64.b64decode(iv)
        tag_bytes = base64.b64decode(tag)

        # Decrypt with GCM
        decryptor = Cipher(
            algorithms.AES(self.key),
            modes.GCM(iv_bytes, tag_bytes),
            backend=default_backend()
        ).decryptor()

        # Add associated data if provided
        if associated_data:
            decryptor.authenticate_additional_data(associated_data.encode())

        plaintext = decryptor.update(ciphertext_bytes) + decryptor.finalize()

        return plaintext.decode()

# Example Usage
def demo_aes_encryption():
    print("=== AES Encryption Demo ===")

    # Initialize cipher
    cipher = AESCipher()

    # Test data
    plaintext = "This is a secret message that needs to be encrypted securely!"

    print(f"Original: {plaintext}")
    print(f"Key (hex): {cipher.key.hex()}")

    # CBC Mode Encryption
    print("\n--- CBC Mode ---")
    cbc_result = cipher.encrypt_cbc(plaintext)
    print(f"Ciphertext (CBC): {cbc_result['ciphertext']}")
    print(f"IV (CBC): {cbc_result['iv']}")

    # CBC Mode Decryption
    decrypted_cbc = cipher.decrypt_cbc(cbc_result['ciphertext'], cbc_result['iv'])
    print(f"Decrypted (CBC): {decrypted_cbc}")

    # GCM Mode Encryption (with authentication)
    print("\n--- GCM Mode (Authenticated) ---")
    associated_data = "user-auth-token-123"
    gcm_result = cipher.encrypt_gcm(plaintext, associated_data)
    print(f"Ciphertext (GCM): {gcm_result['ciphertext']}")
    print(f"IV (GCM): {gcm_result['iv']}")
    print(f"Tag (GCM): {gcm_result['tag']}")
    print(f"Associated Data: {gcm_result['associated_data']}")

    # GCM Mode Decryption
    decrypted_gcm = cipher.decrypt_gcm(
        gcm_result['ciphertext'],
        gcm_result['iv'],
        gcm_result['tag'],
        gcm_result['associated_data']
    )
    print(f"Decrypted (GCM): {decrypted_gcm}")

if __name__ == "__main__":
    demo_aes_encryption()

// ChaCha20-Poly1305 Encryption with Node.js
// Required: npm install tweetnacl

const nacl = require('tweetnacl');
const util = require('tweetnacl-util');

class ChaCha20Cipher {
    constructor() {
        this.keyPair = null;
    }

    // Generate a new 256-bit key
    generateKey() {
        return nacl.randomBytes(32);
    }

    // Generate a 96-bit nonce for ChaCha20
    generateNonce() {
        return nacl.randomBytes(12);
    }

    // Encrypt message with ChaCha20-Poly1305
    encrypt(message, additionalData = null) {
        const key = this.generateKey();
        const nonce = this.generateNonce();

        // Convert message to Uint8Array
        const messageBytes = util.decodeUTF8(message);

        // Add additional authenticated data if provided
        const aadBytes = additionalData ? util.decodeUTF8(additionalData) : null;

        // Encrypt using ChaCha20-Poly1305
        const encrypted = nacl.secretbox(messageBytes, nonce, key);

        if (!encrypted) {
            throw new Error('Encryption failed');
        }

        return {
            ciphertext: util.encodeBase64(encrypted),
            nonce: util.encodeBase64(nonce),
            key: util.encodeBase64(key),
            additionalData: additionalData || null
        };
    }

    // Decrypt ChaCha20-Poly1305 encrypted message
    decrypt(ciphertext, nonce, key, additionalData = null) {
        try {
            const cipherBytes = util.decodeBase64(ciphertext);
            const nonceBytes = util.decodeBase64(nonce);
            const keyBytes = util.decodeBase64(key);

            // Decrypt
            const decrypted = nacl.secretbox.open(cipherBytes, nonceBytes, keyBytes);

            if (!decrypted) {
                throw new Error('Decryption failed - authentication failed');
            }

            return util.encodeUTF8(decrypted);
        } catch (error) {
            throw new Error(`Decryption failed: ${error.message}`);
        }
    }

    // Encrypt with custom key (useful for key exchange scenarios)
    encryptWithKey(message, key, nonce = null) {
        const nonceBytes = nonce ? util.decodeBase64(nonce) : this.generateNonce();
        const messageBytes = util.decodeUTF8(message);
        const keyBytes = typeof key === 'string' ? util.decodeBase64(key) : key;

        const encrypted = nacl.secretbox(messageBytes, nonceBytes, keyBytes);

        if (!encrypted) {
            throw new Error('Encryption failed');
        }

        return {
            ciphertext: util.encodeBase64(encrypted),
            nonce: util.encodeBase64(nonceBytes)
        };
    }

    // Generate key pair for asymmetric encryption using X25519
    generateKeyPair() {
        this.keyPair = nacl.box.keyPair();
        return {
            publicKey: util.encodeBase64(this.keyPair.publicKey),
            secretKey: util.encodeBase64(this.keyPair.secretKey)
        };
    }

    // Derive shared secret using ECDH
    deriveSharedSecret(theirPublicKey, mySecretKey) {
        const theirPubKeyBytes = util.decodeBase64(theirPublicKey);
        const mySecretKeyBytes = util.decodeBase64(mySecretKey);

        const sharedSecret = nacl.box.before(theirPubKeyBytes, mySecretKeyBytes);
        return util.encodeBase64(sharedSecret);
    }
}

// Performance benchmarking
function benchmarkChaCha20() {
    console.log("=== ChaCha20 Performance Benchmark ===");

    const cipher = new ChaCha20Cipher();
    const testSizes = [1024, 10240, 102400, 1024000]; // 1KB, 10KB, 100KB, 1MB

    testSizes.forEach(size => {
        const message = 'A'.repeat(size);

        console.log(`\nTesting message size: ${size} bytes`);

        const startTime = process.hrtime.bigint();

        const encrypted = cipher.encrypt(message);
        const decrypted = cipher.decrypt(
            encrypted.ciphertext,
            encrypted.nonce,
            encrypted.key
        );

        const endTime = process.hrtime.bigint();
        const duration = Number(endTime - startTime) / 1000000; // Convert to milliseconds

        console.log(`Encryption + Decryption time: ${duration.toFixed(2)} ms`);
        console.log(`Throughput: ${(size * 2 / duration * 1000).toFixed(2)} bytes/sec`);
        console.log(`Original matches decrypted: ${message === decrypted}`);
    });
}

// Demo function
function demoChaCha20() {
    console.log("=== ChaCha20-Poly1305 Demo ===");

    const cipher = new ChaCha20Cipher();

    // Basic encryption/decryption
    const message = "This is a secret message encrypted with ChaCha20-Poly1305!";
    const additionalData = "user-session-123";

    console.log(`Original message: ${message}`);
    console.log(`Additional data: ${additionalData}`);

    // Encrypt
    const encrypted = cipher.encrypt(message, additionalData);
    console.log("\nEncrypted:");
    console.log(`Ciphertext: ${encrypted.ciphertext}`);
    console.log(`Nonce: ${encrypted.nonce}`);
    console.log(`Key: ${encrypted.key}`);

    // Decrypt
    const decrypted = cipher.decrypt(
        encrypted.ciphertext,
        encrypted.nonce,
        encrypted.key,
        encrypted.additionalData
    );
    console.log(`\nDecrypted message: ${decrypted}`);
    console.log(`Match with original: ${message === decrypted}`);

    // Key exchange demo
    console.log("\n--- Key Exchange Demo ---");
    const alice = new ChaCha20Cipher();
    const bob = new ChaCha20Cipher();

    const aliceKeys = alice.generateKeyPair();
    const bobKeys = bob.generateKeyPair();

    console.log("Alice's public key:", aliceKeys.publicKey);
    console.log("Bob's public key:", bobKeys.publicKey);

    // Derive shared secrets
    const aliceSharedSecret = alice.deriveSharedSecret(bobKeys.publicKey, aliceKeys.secretKey);
    const bobSharedSecret = bob.deriveSharedSecret(aliceKeys.publicKey, bobKeys.secretKey);

    console.log(`Shared secrets match: ${aliceSharedSecret === bobSharedSecret}`);

    // Use shared secret for encryption
    const secretMessage = "Alice and Bob's secret conversation";
    const aliceEncrypted = alice.encryptWithKey(secretMessage, aliceSharedSecret);
    const bobDecrypted = bob.decryptWithKey(
        aliceEncrypted.ciphertext,
        aliceEncrypted.nonce,
        bobSharedSecret
    );

    console.log(`Secret message decrypted by Bob: ${bobDecrypted}`);
}

// Add decryption method to ChaCha20Cipher class
ChaCha20Cipher.prototype.decryptWithKey = function(ciphertext, nonce, key) {
    return this.decrypt(ciphertext, nonce, key);
};

if (require.main === module) {
    demoChaCha20();
    benchmarkChaCha20();
}

module.exports = ChaCha20Cipher;

// RSA Asymmetric Encryption Implementation
// Required dependencies (Maven):
// <dependency>
//     <groupId>org.bouncycastle</groupId>
//     <artifactId>bcprov-jdk15on</artifactId>
//     <version>1.70</version>
// </dependency>

import java.security.*;
import java.security.spec.*;
import javax.crypto.*;
import javax.crypto.spec.*;
import java.util.Base64;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;

public class RSACipher {

    private KeyPair keyPair;
    private static final String ALGORITHM = "RSA";
    private static final String SIGNATURE_ALGORITHM = "SHA256withRSA";

    public RSACipher(int keySize) throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM);
        keyGen.initialize(keySize);
        this.keyPair = keyGen.generateKeyPair();
    }

    public RSACipher(KeyPair keyPair) {
        this.keyPair = keyPair;
    }

    // Generate RSA key pair
    public static KeyPair generateKeyPair(int keySize) throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM);
        keyGen.initialize(keySize);
        return keyGen.generateKeyPair();
    }

    // Get public key as Base64 string
    public String getPublicKeyBase64() {
        return Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
    }

    // Get private key as Base64 string
    public String getPrivateKeyBase64() {
        return Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
    }

    // Encrypt data with public key
    public String encrypt(String data, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encrypted = cipher.doFinal(data.getBytes());
        return Base64.getEncoder().encodeToString(encrypted);
    }

    // Decrypt data with private key
    public String decrypt(String encryptedData) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        cipher.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
        byte[] decoded = Base64.getDecoder().decode(encryptedData);
        byte[] decrypted = cipher.doFinal(decoded);
        return new String(decrypted);
    }

    // Create digital signature
    public String sign(String data) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
        signature.initSign(keyPair.getPrivate());
        signature.update(data.getBytes());
        byte[] signatureBytes = signature.sign();
        return Base64.getEncoder().encodeToString(signatureBytes);
    }

    // Verify digital signature
    public boolean verify(String data, String signatureStr, PublicKey publicKey) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
        signature.initVerify(publicKey);
        signature.update(data.getBytes());
        byte[] signatureBytes = Base64.getDecoder().decode(signatureStr);
        return signature.verify(signatureBytes);
    }

    // Hybrid encryption: encrypt symmetric key with RSA
    public String encryptLargeData(String data, PublicKey publicKey) throws Exception {
        // Generate random AES key for symmetric encryption
        KeyGenerator aesKeyGen = KeyGenerator.getInstance("AES");
        aesKeyGen.init(256);
        SecretKey aesKey = aesKeyGen.generateKey();

        // Encrypt data with AES
        Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey);
        byte[] iv = new byte[12]; // GCM IV
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
        aesCipher.init(Cipher.ENCRYPT_MODE, aesKey, gcmSpec);

        byte[] encryptedData = aesCipher.doFinal(data.getBytes());

        // Encrypt AES key with RSA
        Cipher rsaCipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        rsaCipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedAesKey = rsaCipher.doFinal(aesKey.getEncoded());

        // Combine IV, encrypted AES key, and encrypted data
        byte[] combined = new byte[iv.length + encryptedAesKey.length + encryptedData.length];
        System.arraycopy(iv, 0, combined, 0, iv.length);
        System.arraycopy(encryptedAesKey, 0, combined, iv.length, encryptedAesKey.length);
        System.arraycopy(encryptedData, 0, combined, iv.length + encryptedAesKey.length, encryptedData.length);

        return Base64.getEncoder().encodeToString(combined);
    }

    // Hybrid decryption
    public String decryptLargeData(String combinedData) throws Exception {
        byte[] combined = Base64.getDecoder().decode(combinedData);

        // Extract IV, encrypted AES key, and encrypted data
        byte[] iv = new byte[12];
        int rsaKeySize = ((RSAPrivateKey) keyPair.getPrivate()).getModulus().bitLength() / 8;
        byte[] encryptedAesKey = new byte[rsaKeySize];
        byte[] encryptedData = new byte[combined.length - iv.length - encryptedAesKey.length];

        System.arraycopy(combined, 0, iv, 0, iv.length);
        System.arraycopy(combined, iv.length, encryptedAesKey, 0, encryptedAesKey.length);
        System.arraycopy(combined, iv.length + encryptedAesKey.length, encryptedData, 0, encryptedData.length);

        // Decrypt AES key with RSA
        Cipher rsaCipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        rsaCipher.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
        byte[] aesKeyBytes = rsaCipher.doFinal(encryptedAesKey);
        SecretKey aesKey = new SecretKeySpec(aesKeyBytes, "AES");

        // Decrypt data with AES
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
        Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
        aesCipher.init(Cipher.DECRYPT_MODE, aesKey, gcmSpec);
        byte[] decryptedData = aesCipher.doFinal(encryptedData);

        return new String(decryptedData);
    }

    public static void main(String[] args) {
        try {
            System.out.println("=== RSA Asymmetric Encryption Demo ===");

            // Generate RSA key pair
            RSACipher rsaCipher = new RSACipher(2048);

            String message = "This is a secret message that will be encrypted with RSA!";
            System.out.println("Original message: " + message);

            // Get keys as strings
            String publicKeyStr = rsaCipher.getPublicKeyBase64();
            String privateKeyStr = rsaCipher.getPrivateKeyBase64();

            System.out.println("\nPublic Key (Base64): " + publicKeyStr.substring(0, 50) + "...");
            System.out.println("Private Key (Base64): " + privateKeyStr.substring(0, 50) + "...");

            // Encrypt with public key
            String encrypted = rsaCipher.encrypt(message, rsaCipher.keyPair.getPublic());
            System.out.println("\nEncrypted message: " + encrypted);

            // Decrypt with private key
            String decrypted = rsaCipher.decrypt(encrypted);
            System.out.println("Decrypted message: " + decrypted);
            System.out.println("Original matches decrypted: " + message.equals(decrypted));

            // Digital signature demo
            System.out.println("\n--- Digital Signature Demo ---");
            String document = "Important document to be signed";
            String signature = rsaCipher.sign(document);
            System.out.println("Document: " + document);
            System.out.println("Signature: " + signature.substring(0, 50) + "...");

            boolean isValid = rsaCipher.verify(document, signature, rsaCipher.keyPair.getPublic());
            System.out.println("Signature is valid: " + isValid);

            // Try with modified document
            boolean isInvalid = rsaCipher.verify(document + " modified", signature, rsaCipher.keyPair.getPublic());
            System.out.println("Modified document signature is valid: " + isInvalid);

            // Hybrid encryption demo for large data
            System.out.println("\n--- Hybrid Encryption Demo (Large Data) ---");
            String largeMessage = "This is a very long message that exceeds RSA's data size limitations. ".repeat(10);
            System.out.println("Original message length: " + largeMessage.length() + " characters");

            String hybridEncrypted = rsaCipher.encryptLargeData(largeMessage, rsaCipher.keyPair.getPublic());
            System.out.println("Hybrid encrypted length: " + hybridEncrypted.length());

            String hybridDecrypted = rsaCipher.decryptLargeData(hybridEncrypted);
            System.out.println("Hybrid decrypted length: " + hybridDecrypted.length() + " characters");
            System.out.println("Large message matches hybrid decrypted: " + largeMessage.equals(hybridDecrypted));

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

// Elliptic Curve Cryptography Implementation in Go
// Run with: go run ecc.go

package main

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"time"
)

// ECDSAKeyPair represents an ECDSA key pair
type ECDSAKeyPair struct {
	PrivateKey *ecdsa.PrivateKey
	PublicKey  *ecdsa.PublicKey
}

// Ed25519KeyPair represents an Ed25519 key pair
type Ed25519KeyPair struct {
	PrivateKey ed25519.PrivateKey
	PublicKey  ed25519.PublicKey
}

// ECDHSecret represents a shared secret from ECDH
type ECDHSecret struct {
	Secret []byte
	X      *big.Int
	Y      *big.Int
}

// Generate ECDSA key pair for specified curve
func GenerateECDSAKeyPair(curve elliptic.Curve) (*ECDSAKeyPair, error) {
	privateKey, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate ECDSA key pair: %v", err)
	}

	return &ECDSAKeyPair{
		PrivateKey: privateKey,
		PublicKey:  &privateKey.PublicKey,
	}, nil
}

// Generate Ed25519 key pair
func GenerateEd25519KeyPair() (*Ed25519KeyPair, error) {
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate Ed25519 key pair: %v", err)
	}

	return &Ed25519KeyPair{
		PrivateKey: privateKey,
		PublicKey:  publicKey,
	}, nil
}

// Sign data with ECDSA
func (kp *ECDSAKeyPair) Sign(data []byte) (string, string, error) {
	hash := sha256.Sum256(data)

	r, s, err := ecdsa.Sign(rand.Reader, kp.PrivateKey, hash[:])
	if err != nil {
		return "", "", fmt.Errorf("failed to sign data: %v", err)
	}

	signature := append(r.Bytes(), s.Bytes()...)

	return base64.StdEncoding.EncodeToString(signature),
		fmt.Sprintf("%032x", r) + fmt.Sprintf("%032x", s),
		nil
}

// Verify ECDSA signature
func (kp *ECDSAKeyPair) Verify(data []byte, signatureBase64 string) bool {
	signature, err := base64.StdEncoding.DecodeString(signatureBase64)
	if err != nil {
		return false
	}

	if len(signature) != 64 { // 32 bytes for r, 32 bytes for s
		return false
	}

	r := new(big.Int).SetBytes(signature[:32])
	s := new(big.Int).SetBytes(signature[32:])

	hash := sha256.Sum256(data)
	return ecdsa.Verify(kp.PublicKey, hash[:], r, s)
}

// Sign data with Ed25519
func (kp *Ed25519KeyPair) Sign(data []byte) string {
	signature := ed25519.Sign(kp.PrivateKey, data)
	return base64.StdEncoding.EncodeToString(signature)
}

// Verify Ed25519 signature
func (kp *Ed25519KeyPair) Verify(data []byte, signatureBase64 string) bool {
	signature, err := base64.StdEncoding.DecodeString(signatureBase64)
	if err != nil {
		return false
	}
	return ed25519.Verify(kp.PublicKey, data, signature)
}

// Perform ECDH key exchange
func (kp *ECDSAKeyPair) ECDH(theirPublicKey *ecdsa.PublicKey) (*ECDHSecret, error) {
	// Calculate shared secret
	scalarX, _ := theirPublicKey.Curve.ScalarMult(theirPublicKey.X, theirPublicKey.Y, kp.PrivateKey.D.Bytes())

	if scalarX == nil {
		return nil, fmt.Errorf("failed to compute ECDH secret")
	}

	// Hash the x-coordinate to get a uniform secret
	hash := sha256.Sum256(scalarX.Bytes())

	return &ECDHSecret{
		Secret: hash[:],
		X:      scalarX,
		Y:      theirPublicKey.Y,
	}, nil
}

// Export ECDSA public key to PEM format
func (kp *ECDSAKeyPair) ExportPublicKeyPEM() (string, error) {
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(kp.PublicKey)
	if err != nil {
		return "", fmt.Errorf("failed to marshal public key: %v", err)
	}

	publicKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyBytes,
	})

	return string(publicKeyPEM), nil
}

// Export Ed25519 public key to base64
func (kp *Ed25519KeyPair) ExportPublicKeyBase64() string {
	return base64.StdEncoding.EncodeToString(kp.PublicKey)
}

// Get ECDSA key coordinates
func (kp *ECDSAKeyPair) GetCoordinates() (string, string) {
	return kp.PublicKey.X.String(), kp.PublicKey.Y.String()
}

// Benchmark different curves
func benchmarkCurves() {
	fmt.Println("=== ECC Curve Performance Benchmark ===")

	curves := map[string]elliptic.Curve{
		"P-256":  elliptic.P256(),
		"P-384":  elliptic.P384(),
		"P-521":  elliptic.P521(),
		"secp256k1": nil, // Will be handled separately
	}

	for name, curve := range curves {
		fmt.Printf("\n--- %s ---\n", name)

		var kp *ECDSAKeyPair
		var err error

		if name == "secp256k1" {
			// For secp256k1, we'd need an external library like btcec
			fmt.Println("secp256k1 requires external library (btcec)")
			continue
		}

		kp, err = GenerateECDSAKeyPair(curve)
		if err != nil {
			fmt.Printf("Failed to generate key pair: %v\n", err)
			continue
		}

		// Test signing performance
		data := []byte("Benchmark data for ECC performance testing")
		iterations := 1000

		start := makeTimestamp()
		for i := 0; i < iterations; i++ {
			_, _, err := kp.Sign(data)
			if err != nil {
				fmt.Printf("Sign failed: %v\n", err)
				break
			}
		}
		signDuration := makeTimestamp() - start

		// Test verification performance
		signature, _, err := kp.Sign(data)
		if err != nil {
			fmt.Printf("Failed to create signature for verification test\n")
			continue
		}

		start = makeTimestamp()
		for i := 0; i < iterations; i++ {
			kp.Verify(data, signature)
		}
		verifyDuration := makeTimestamp() - start

		fmt.Printf("Key size: %d bits\n", curve.Params().BitSize)
		fmt.Printf("Signing: %d ops in %.2f ms (%.2f ops/sec)\n",
			iterations, float64(signDuration)/1000000,
			float64(iterations)/float64(signDuration)*1000000000)
		fmt.Printf("Verification: %d ops in %.2f ms (%.2f ops/sec)\n",
			iterations, float64(verifyDuration)/1000000,
			float64(iterations)/float64(verifyDuration)*1000000000)
	}
}

// Helper function for timestamp
func makeTimestamp() int64 {
	return time.Now().UnixNano()
}

// Main demo function
func main() {
	fmt.Println("=== Elliptic Curve Cryptography Demo ===")

	// ECDSA Demo with P-256 curve
	fmt.Println("\n--- ECDSA with P-256 ---")
	ecdsaKp, err := GenerateECDSAKeyPair(elliptic.P256())
	if err != nil {
		log.Fatalf("Failed to generate ECDSA key pair: %v", err)
	}

	message := []byte("This message will be signed with ECDSA")
	fmt.Printf("Original message: %s\n", string(message))

	// Sign message
	signature, sigHex, err := ecdsaKp.Sign(message)
	if err != nil {
		log.Fatalf("Failed to sign message: %v", err)
	}

	fmt.Printf("Signature (Base64): %s\n", signature[:64]+"...")
	fmt.Printf("Signature (Hex): %s\n", sigHex)

	// Verify signature
	isValid := ecdsaKp.Verify(message, signature)
	fmt.Printf("Signature verification: %t\n", isValid)

	// Try with modified message
	modifiedMessage := []byte("This message will be signed with ECDSA - MODIFIED")
	isValid = ecdsaKp.Verify(modifiedMessage, signature)
	fmt.Printf("Modified message verification: %t\n", isValid)

	// Export public key
	pubKeyPEM, err := ecdsaKp.ExportPublicKeyPEM()
	if err != nil {
		log.Fatalf("Failed to export public key: %v", err)
	}
	fmt.Printf("Public Key (PEM format):\n%s\n", pubKeyPEM)

	// Get coordinates
	x, y := ecdsaKp.GetCoordinates()
	fmt.Printf("Public Key Coordinates:\n")
	fmt.Printf("  X: %s\n", x)
	fmt.Printf("  Y: %s\n", y)

	// Ed25519 Demo
	fmt.Println("\n--- Ed25519 ---")
	edKp, err := GenerateEd25519KeyPair()
	if err != nil {
		log.Fatalf("Failed to generate Ed25519 key pair: %v", err)
	}

	edMessage := []byte("This message will be signed with Ed25519")
	edSignature := edKp.Sign(edMessage)

	fmt.Printf("Original message: %s\n", string(edMessage))
	fmt.Printf("Signature (Base64): %s\n", edSignature[:64]+"...")

	edIsValid := edKp.Verify(edMessage, edSignature)
	fmt.Printf("Signature verification: %t\n", edIsValid)

	// ECDH Key Exchange Demo
	fmt.Println("\n--- ECDH Key Exchange ---")

	// Alice's key pair
	aliceKp, err := GenerateECDSAKeyPair(elliptic.P256())
	if err != nil {
		log.Fatalf("Failed to generate Alice's key pair: %v", err)
	}

	// Bob's key pair
	bobKp, err := GenerateECDSAKeyPair(elliptic.P256())
	if err != nil {
		log.Fatalf("Failed to generate Bob's key pair: %v", err)
	}

	// Alice computes shared secret using Bob's public key
	aliceSecret, err := aliceKp.ECDH(bobKp.PublicKey)
	if err != nil {
		log.Fatalf("Alice failed to compute shared secret: %v", err)
	}

	// Bob computes shared secret using Alice's public key
	bobSecret, err := bobKp.ECDH(aliceKp.PublicKey)
	if err != nil {
		log.Fatalf("Bob failed to compute shared secret: %v", err)
	}

	fmt.Printf("Alice's shared secret: %s\n", hex.EncodeToString(aliceSecret.Secret))
	fmt.Printf("Bob's shared secret: %s\n", hex.EncodeToString(bobSecret.Secret))
	fmt.Printf("Shared secrets match: %t\n",
		hex.EncodeToString(aliceSecret.Secret) == hex.EncodeToString(bobSecret.Secret))

	// Run benchmarks
	benchmarkCurves()

	fmt.Println("\n=== ECC Demo Complete ===")
}

# Hash Functions and Digital Signatures
# Required: pip install cryptography pycryptodome

import hashlib
import hmac
import json
import time
from typing import List, Dict, Any, Optional
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import utils
from cryptography import x509
from cryptography.x509.oid import NameOID
import base64
import os
from datetime import datetime, timedelta

class HashFunctionManager:
    """Comprehensive hash function utilities"""

    @staticmethod
    def sha256(data: str) -> str:
        """SHA-256 hash function"""
        return hashlib.sha256(data.encode()).hexdigest()

    @staticmethod
    def sha512(data: str) -> str:
        """SHA-512 hash function"""
        return hashlib.sha512(data.encode()).hexdigest()

    @staticmethod
    def sha3_256(data: str) -> str:
        """SHA3-256 hash function"""
        return hashlib.sha3_256(data.encode()).hexdigest()

    @staticmethod
    def md5(data: str) -> str:
        """MD5 hash function (for legacy compatibility only)"""
        return hashlib.md5(data.encode()).hexdigest()

    @staticmethod
    def blake2b(data: str) -> str:
        """BLAKE2b hash function"""
        return hashlib.blake2b(data.encode()).hexdigest()

    @staticmethod
    def hmac_sha256(data: str, key: str) -> str:
        """HMAC with SHA-256"""
        return hmac.new(key.encode(), data.encode(), hashlib.sha256).hexdigest()

    @staticmethod
    def file_hash(file_path: str, algorithm: str = 'sha256') -> str:
        """Calculate hash of a file"""
        hash_func = getattr(hashlib, algorithm)()

        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_func.update(chunk)

        return hash_func.hexdigest()

class MerkleTree:
    """Merkle Tree implementation for data integrity verification"""

    def __init__(self, data_chunks: List[str]):
        self.leaves = [HashFunctionManager.sha256(chunk) for chunk in data_chunks]
        self.tree = self._build_tree(self.leaves)
        self.root = self.tree[0] if self.tree else ""

    def _build_tree(self, nodes: List[str]) -> List[str]:
        """Build merkle tree from leaf nodes"""
        if len(nodes) == 1:
            return nodes

        tree_level = []

        # Process pairs of nodes
        for i in range(0, len(nodes), 2):
            if i + 1 < len(nodes):
                # Hash concatenation of two nodes
                combined = nodes[i] + nodes[i + 1]
                tree_level.append(HashFunctionManager.sha256(combined))
            else:
                # Handle odd number of nodes by duplicating the last one
                combined = nodes[i] + nodes[i]
                tree_level.append(HashFunctionManager.sha256(combined))

        # Recursively build tree
        return self._build_tree(tree_level)

    def get_proof(self, data_chunk: str) -> List[Dict[str, Any]]:
        """Generate Merkle proof for a data chunk"""
        chunk_hash = HashFunctionManager.sha256(data_chunk)

        if chunk_hash not in self.leaves:
            raise ValueError("Data chunk not found in tree")

        proof = []
        current_level = self.leaves
        current_hash = chunk_hash
        index = current_level.index(current_hash)

        while len(current_level) > 1:
            # Determine if node is left or right child
            is_right_child = index % 2 == 1
            sibling_index = index - 1 if is_right_child else index + 1

            if sibling_index < len(current_level):
                proof.append({
                    'hash': current_level[sibling_index],
                    'direction': 'left' if is_right_child else 'right'
                })

            # Move to next level
            index = index // 2
            current_level = self._build_tree(current_level)

        return proof

    def verify_proof(self, data_chunk: str, proof: List[Dict[str, Any]]) -> bool:
        """Verify Merkle proof for a data chunk"""
        current_hash = HashFunctionManager.sha256(data_chunk)

        for proof_element in proof:
            if proof_element['direction'] == 'left':
                combined = proof_element['hash'] + current_hash
            else:
                combined = current_hash + proof_element['hash']

            current_hash = HashFunctionManager.sha256(combined)

        return current_hash == self.root

class DigitalSignatureManager:
    """Digital signature management with RSA and certificates"""

    def __init__(self):
        self.private_key = None
        self.public_key = None
        self.certificate = None

    def generate_rsa_key_pair(self, key_size: int = 2048) -> None:
        """Generate RSA key pair for signing"""
        self.private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=key_size
        )
        self.public_key = self.private_key.public_key()

    def create_self_signed_certificate(self, common_name: str = "Test Certificate") -> None:
        """Create a self-signed X.509 certificate"""
        if not self.private_key:
            raise ValueError("Private key not generated")

        subject = issuer = x509.Name([
            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
            x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Test Org"),
            x509.NameAttribute(NameOID.COMMON_NAME, common_name),
        ])

        self.certificate = (
            x509.CertificateBuilder()
            .subject_name(subject)
            .issuer_name(issuer)
            .public_key(self.public_key)
            .serial_number(x509.random_serial_number())
            .not_valid_before(datetime.utcnow())
            .not_valid_after(datetime.utcnow() + timedelta(days=365))
            .add_extension(
                x509.SubjectAlternativeName([
                    x509.DNSName(common_name),
                ]),
                critical=False,
            )
            .sign(self.private_key, hashes.SHA256())
        )

    def sign_message(self, message: str) -> str:
        """Sign a message with RSA private key"""
        if not self.private_key:
            raise ValueError("Private key not generated")

        signature = self.private_key.sign(
            message.encode(),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )

        return base64.b64encode(signature).decode()

    def verify_signature(self, message: str, signature_b64: str, public_key_pem: str = None) -> bool:
        """Verify a digital signature"""
        try:
            signature = base64.b64decode(signature_b64)

            # Use provided public key or own public key
            if public_key_pem:
                pub_key = serialization.load_pem_public_key(
                    public_key_pem.encode()
                )
            else:
                pub_key = self.public_key

            pub_key.verify(
                signature,
                message.encode(),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )

            return True
        except Exception:
            return False

    def export_public_key_pem(self) -> str:
        """Export public key in PEM format"""
        if not self.public_key:
            raise ValueError("Public key not available")

        pem = self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )

        return pem.decode()

    def export_certificate_pem(self) -> str:
        """Export certificate in PEM format"""
        if not self.certificate:
            raise ValueError("Certificate not created")

        pem = self.certificate.public_bytes(serialization.Encoding.PEM)
        return pem.decode()

class BlockchainHashing:
    """Blockchain-style hashing with proof of work"""

    @staticmethod
    def hash_block(block_data: Dict[str, Any], previous_hash: str, nonce: int = 0) -> str:
        """Hash a block with mining capability"""
        block_string = json.dumps({
            'data': block_data,
            'previous_hash': previous_hash,
            'nonce': nonce,
            'timestamp': time.time()
        }, sort_keys=True)

        return HashFunctionManager.sha256(block_string)

    @staticmethod
    def mine_block(block_data: Dict[str, Any], previous_hash: str, difficulty: int = 4) -> Dict[str, Any]:
        """Mine a block with proof of work"""
        target = "0" * difficulty
        nonce = 0

        while True:
            block_hash = BlockchainHashing.hash_block(block_data, previous_hash, nonce)

            if block_hash.startswith(target):
                return {
                    'data': block_data,
                    'previous_hash': previous_hash,
                    'nonce': nonce,
                    'hash': block_hash,
                    'timestamp': time.time(),
                    'difficulty': difficulty
                }

            nonce += 1

    @staticmethod
    def verify_chain(chain: List[Dict[str, Any]]) -> bool:
        """Verify blockchain integrity"""
        for i in range(1, len(chain)):
            current_block = chain[i]
            previous_block = chain[i - 1]

            # Verify hash
            expected_hash = BlockchainHashing.hash_block(
                current_block['data'],
                current_block['previous_hash'],
                current_block['nonce']
            )

            if expected_hash != current_block['hash']:
                return False

            # Verify previous hash link
            if current_block['previous_hash'] != previous_block['hash']:
                return False

        return True

def demo_hash_functions():
    print("=== Hash Functions Demo ===")

    data = "Hello, Cryptography!"

    print(f"Original data: {data}")
    print(f"SHA-256: {HashFunctionManager.sha256(data)}")
    print(f"SHA-512: {HashFunctionManager.sha512(data)}")
    print(f"SHA3-256: {HashFunctionManager.sha3_256(data)}")
    print(f"MD5: {HashFunctionManager.md5(data)}")
    print(f"BLAKE2b: {HashFunctionManager.blake2b(data)}")

    # HMAC Demo
    key = "secret-key-123"
    hmac_result = HashFunctionManager.hmac_sha256(data, key)
    print(f"\nHMAC-SHA256 with key '{key}': {hmac_result}")

def demo_merkle_tree():
    print("\n=== Merkle Tree Demo ===")

    # Sample data chunks
    documents = [
        "Document 1: Important contract",
        "Document 2: Financial report",
        "Document 3: Medical records",
        "Document 4: Legal notice",
        "Document 5: Tax information"
    ]

    merkle_tree = MerkleTree(documents)

    print("Documents:")
    for i, doc in enumerate(documents):
        print(f"  {i+1}. {doc}")

    print(f"\nMerkle Root: {merkle_tree.root}")

    # Generate and verify proof for document 3
    target_doc = documents[2]
    proof = merkle_tree.get_proof(target_doc)

    print(f"\nProof for document '{target_doc[:20]}...':")
    for i, proof_element in enumerate(proof):
        print(f"  Step {i+1}: {proof_element['hash'][:16]}... ({proof_element['direction']})")

    # Verify proof
    is_valid = merkle_tree.verify_proof(target_doc, proof)
    print(f"\nProof verification: {is_valid}")

    # Try with modified document
    modified_doc = target_doc + " (modified)"
    is_valid_modified = merkle_tree.verify_proof(modified_doc, proof)
    print(f"Modified document verification: {is_valid_modified}")

def demo_digital_signatures():
    print("\n=== Digital Signatures Demo ===")

    # Initialize signature manager
    sig_manager = DigitalSignatureManager()
    sig_manager.generate_rsa_key_pair(2048)
    sig_manager.create_self_signed_certificate("example.com")

    # Sign a message
    message = "This is a legally binding electronic document"
    signature = sig_manager.sign_message(message)

    print(f"Original message: {message}")
    print(f"Signature: {signature[:64]}...")

    # Export keys and certificate
    public_key_pem = sig_manager.export_public_key_pem()
    certificate_pem = sig_manager.export_certificate_pem()

    print(f"\nPublic Key (PEM format):")
    print(public_key_pem[:100] + "...")

    print(f"\nCertificate (PEM format):")
    print(certificate_pem[:100] + "...")

    # Verify signature
    is_valid = sig_manager.verify_signature(message, signature)
    print(f"\nSignature verification: {is_valid}")

    # Try with modified message
    modified_message = message + " (modified)"
    is_valid_modified = sig_manager.verify_signature(modified_message, signature)
    print(f"Modified message verification: {is_valid_modified}")

    # Verify with exported public key
    is_valid_exported = sig_manager.verify_signature(message, signature, public_key_pem)
    print(f"Verification with exported public key: {is_valid_exported}")

def demo_blockchain_hashing():
    print("\n=== Blockchain Hashing Demo ===")

    # Create a simple blockchain
    block1 = BlockchainHashing.mine_block(
        {"transactions": ["Alice pays Bob 10 BTC"]},
        "0" * 64,  # Genesis block previous hash
        difficulty=2
    )

    block2 = BlockchainHashing.mine_block(
        {"transactions": ["Bob pays Charlie 5 BTC"]},
        block1['hash'],
        difficulty=2
    )

    block3 = BlockchainHashing.mine_block(
        {"transactions": ["Charlie pays Dave 3 BTC"]},
        block2['hash'],
        difficulty=2
    )

    blockchain = [block1, block2, block3]

    print("Blockchain:")
    for i, block in enumerate(blockchain):
        print(f"\nBlock {i+1}:")
        print(f"  Data: {block['data']}")
        print(f"  Nonce: {block['nonce']}")
        print(f"  Hash: {block['hash']}")
        print(f"  Previous Hash: {block['previous_hash']}")

    # Verify blockchain integrity
    is_valid_chain = BlockchainHashing.verify_chain(blockchain)
    print(f"\nBlockchain integrity: {is_valid_chain}")

    # Tamper with blockchain and verify
    tampered_blockchain = [block1, block2, block3.copy()]
    tampered_blockchain[2]['data'] = {"transactions": ["Charlie pays Eve 100 BTC"]}

    is_tampered_valid = BlockchainHashing.verify_chain(tampered_blockchain)
    print(f"Tampered blockchain integrity: {is_tampered_valid}")

def main():
    print("=== Cryptography: Hash Functions and Digital Signatures ===")

    demo_hash_functions()
    demo_merkle_tree()
    demo_digital_signatures()
    demo_blockchain_hashing()

    print("\n=== Demo Complete ===")

if __name__ == "__main__":
    main()

// Key Exchange Protocols Implementation in Rust
// Cargo.toml dependencies:
// [dependencies]
// ring = "0.16"
// x25519-dalek = "2.0"
// curve25519-dalek = "4.0"
// sha2 = "0.10"
// rand = "0.8"
// serde = { version = "1.0", features = ["derive"] }
// base64 = "0.21"

use curve25519_dalek::{ristretto::RistrettoPoint, scalar::Scalar, constants::RISTRETTO_BASEPOINT_POINT};
use ring::{rand::SecureRandom, agreement, digest};
use sha2::{Sha256, Digest};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct DiffieHellmanKeyPair {
    pub private_key: Vec<u8>,
    pub public_key: Vec<u8>,
}

#[derive(Debug, Clone)]
pub struct ECDHKeyPair {
    pub private_key: Scalar,
    pub public_key: RistrettoPoint,
}

#[derive(Debug, Clone)]
pub struct KeyExchangeResult {
    pub shared_secret: Vec<u8>,
    pub session_key: Vec<u8>,
    pub exchange_id: String,
}

pub struct KeyExchangeManager {
    rng: ring::rand::SystemRandom,
}

impl KeyExchangeManager {
    pub fn new() -> Self {
        Self {
            rng: ring::rand::SystemRandom::new(),
        }
    }

    /// Traditional Diffie-Hellman key exchange
    pub fn diffie_hellman_generate(&self, key_size: usize) -> Result<DiffieHellmanKeyPair, &'static str> {
        let mut private_key = vec![0u8; key_size];
        let mut public_key = vec![0u8; key_size];

        // Generate private key
        self.rng.fill(&mut private_key);

        // For demonstration, we'll use a simple generator (2) and large prime
        // In practice, use well-known safe primes like RFC 3526 groups
        let generator = 2u64;
        let prime = 0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381FFFFFFFFFFFFFFFFu64;

        // Calculate public key: g^private_key mod p
        let private_val = u64::from_le_bytes([
            private_key[0], private_key[1], private_key[2], private_key[3],
            private_key[4], private_key[5], private_key[6], private_key[7],
        ]);

        let public_val = generator.modpow(private_val, prime);
        public_key[..8].copy_from_slice(&public_val.to_le_bytes());

        Ok(DiffieHellmanKeyPair {
            private_key,
            public_key,
        })
    }

    /// Compute shared secret using Diffie-Hellman
    pub fn diffie_hellman_compute_secret(
        &self,
        my_private: &[u8],
        their_public: &[u8],
        key_size: usize,
    ) -> Result<Vec<u8>, &'static str> {
        // Extract values (simplified for demonstration)
        let private_val = u64::from_le_bytes([
            my_private[0], my_private[1], my_private[2], my_private[3],
            my_private[4], my_private[5], my_private[6], my_private[7],
        ]);

        let their_public_val = u64::from_le_bytes([
            their_public[0], their_public[1], their_public[2], their_public[3],
            their_public[4], their_public[5], their_public[6], their_public[7],
        ]);

        let prime = 0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381FFFFFFFFFFFFFFFFu64;

        // Compute shared secret: their_public^my_private mod p
        let shared_val = their_public_val.modpow(private_val, prime);

        let mut shared_secret = vec![0u8; key_size];
        shared_secret[..8].copy_from_slice(&shared_val.to_le_bytes());

        Ok(shared_secret)
    }

    /// Generate ECDH key pair using Ristretto (Curve25519 variant)
    pub fn ecdh_generate(&self) -> Result<ECDHKeyPair, &'static str> {
        let mut private_bytes = [0u8; 32];
        self.rng.fill(&mut private_bytes);

        let private_key = Scalar::from_bytes_mod_order(private_bytes);
        let public_key = RistrettoPoint::mult_base_private(&private_key);

        Ok(ECDHKeyPair {
            private_key,
            public_key,
        })
    }

    /// Compute ECDH shared secret
    pub fn ecdh_compute_secret(
        &self,
        my_private: &Scalar,
        their_public: &RistrettoPoint,
    ) -> Vec<u8> {
        let shared_point = their_public * my_private;
        shared_point.compress().to_bytes().to_vec()
    }

    /// Post-quantum key resistant key exchange (Kyber-like)
    pub fn post_quantum_key_exchange(&self) -> Result<(Vec<u8>, Vec<u8>), &'static str> {
        // Simplified post-quantum key exchange based on lattice-based cryptography
        // In practice, use libraries like pqcrypto or ML-KEM

        // Generate random lattice parameters
        let mut seed = [0u8; 32];
        self.rng.fill(&mut seed);

        // Generate "public key" (commitment) and "private key" (trapdoor)
        let mut public_key = [0u8; 1024];
        let mut private_key = [0u8; 2048];

        // Fill with pseudo-random data based on seed
        let mut hasher = Sha256::new();
        hasher.update(seed);
        let hash = hasher.finalize();

        for i in 0..1024 {
            public_key[i] = hash[i % 32] ^ (i as u8);
        }

        for i in 0..2048 {
            private_key[i] = hash[i % 32] ^ ((i >> 3) as u8);
        }

        Ok((public_key.to_vec(), private_key.to_vec()))
    }

    /// Derive session key from shared secret
    pub fn derive_session_key(
        &self,
        shared_secret: &[u8],
        context: &[u8],
        key_length: usize,
    ) -> Vec<u8> {
        let mut hasher = Sha256::new();
        hasher.update(shared_secret);
        hasher.update(context);

        let mut key = vec![0u8; key_length];
        let hash = hasher.finalize();

        // Simple key derivation: repeat hash if needed
        for i in 0..key_length {
            key[i] = hash[i % 32];
        }

        key
    }

    /// Authenticated key exchange using digital signatures
    pub fn authenticated_key_exchange(
        &self,
        identity: &str,
    ) -> Result<(Vec<u8>, Vec<u8>), &'static str> {
        // Generate ECDH key pair
        let ecdh_keypair = self.ecdh_generate()?;

        // Create identity-based signature
        let mut hasher = Sha256::new();
        hasher.update(identity.as_bytes());
        hasher.update(ecdh_keypair.public_key.compress().to_bytes());
        let signature = hasher.finalize();

        // Return public key and signature
        Ok((
            ecdh_keypair.public_key.compress().to_bytes().to_vec(),
            signature.to_vec(),
        ))
    }

    /// Verify authenticated key exchange
    pub fn verify_authenticated_exchange(
        &self,
        identity: &str,
        public_key_bytes: &[u8],
        signature: &[u8],
    ) -> bool {
        let mut hasher = Sha256::new();
        hasher.update(identity.as_bytes());
        hasher.update(public_key_bytes);
        let expected_signature = hasher.finalize();

        signature == expected_signature.as_slice()
    }

    /// Perfect Forward Secrecy with ephemeral keys
    pub fn perfect_forward_secrecy_exchange(&self) -> Result<KeyExchangeResult, &'static str> {
        // Generate ephemeral ECDH key pair
        let ephemeral_keypair = self.ecdh_generate()?;

        // Generate long-term key pair
        let longterm_keypair = self.ecdh_generate()?;

        // Compute shared secrets
        let shared_secret = self.ecdh_compute_secret(&ephemeral_keypair.private_key, &longterm_keypair.public_key);

        // Derive session keys
        let session_key = self.derive_session_key(&shared_secret, b"PFS-session", 32);

        // Generate exchange ID
        let exchange_id = format!("pfs-{}", std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs());

        Ok(KeyExchangeResult {
            shared_secret,
            session_key,
            exchange_id,
        })
    }

    /// Group key exchange (TreeKEM-like protocol)
    pub fn group_key_exchange(&self, participants: &[String]) -> Result<HashMap<String, Vec<u8>>, &'static str> {
        let mut group_keys = HashMap::new();

        // Sort participants for deterministic ordering
        let mut sorted_participants = participants.clone();
        sorted_participants.sort();

        // Generate a master secret
        let mut master_secret = [0u8; 32];
        self.rng.fill(&mut master_secret);

        // Derive individual keys for each participant
        for (i, participant) in sorted_participants.iter().enumerate() {
            let mut hasher = Sha256::new();
            hasher.update(&master_secret);
            hasher.update(participant.as_bytes());
            hasher.update(&(i as u64).to_le_bytes());

            let participant_key = hasher.finalize();
            group_keys.insert(participant.clone(), participant_key.to_vec());
        }

        Ok(group_keys)
    }

    /// Key rotation for long-term security
    pub fn rotate_key(&self, current_key: &[u8], rotation_id: &str) -> Vec<u8> {
        let mut hasher = Sha256::new();
        hasher.update(current_key);
        hasher.update(rotation_id.as_bytes());
        hasher.update(&std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
            .to_le_bytes());

        hasher.finalize().to_vec()
    }
}

fn benchmark_key_exchanges() {
    println!("=== Key Exchange Performance Benchmark ===");

    let manager = KeyExchangeManager::new();
    let iterations = 1000;

    // Benchmark Diffie-Hellman
    let start = std::time::Instant::now();
    for _ in 0..iterations {
        let _ = manager.diffie_hellman_generate(256);
    }
    let dh_duration = start.elapsed();

    // Benchmark ECDH
    let start = std::time::Instant::now();
    for _ in 0..iterations {
        let _ = manager.ecdh_generate();
    }
    let ecdh_duration = start.elapsed();

    println!("Diffie-Hellman ({} ops): {:?}", iterations, dh_duration);
    println!("ECDH ({} ops): {:?}", iterations, ecdh_duration);

    let dh_ops_per_sec = iterations as f64 / dh_duration.as_secs_f64();
    let ecdh_ops_per_sec = iterations as f64 / ecdh_duration.as_secs_f64();

    println!("Diffie-Hellman: {:.2} ops/sec", dh_ops_per_sec);
    println!("ECDH: {:.2} ops/sec", ecdh_ops_per_sec);
}

fn main() {
    println!("=== Key Exchange Protocols Demo ===");

    let manager = KeyExchangeManager::new();

    // Traditional Diffie-Hellman Demo
    println!("\n--- Traditional Diffie-Hellman ---");
    let alice_dh = manager.diffie_hellman_generate(256).unwrap();
    let bob_dh = manager.diffie_hellman_generate(256).unwrap();

    println!("Alice's private key: {}", hex::encode(&alice_dh.private_key[..16]));
    println!("Alice's public key: {}", hex::encode(&alice_dh.public_key[..16]));
    println!("Bob's private key: {}", hex::encode(&bob_dh.private_key[..16]));
    println!("Bob's public key: {}", hex::encode(&bob_dh.public_key[..16]));

    let alice_shared = manager.diffie_hellman_compute_secret(
        &alice_dh.private_key,
        &bob_dh.public_key,
        256,
    ).unwrap();

    let bob_shared = manager.diffie_hellman_compute_secret(
        &bob_dh.private_key,
        &alice_dh.public_key,
        256,
    ).unwrap();

    println!("Alice's shared secret: {}", hex::encode(&alice_shared[..16]));
    println!("Bob's shared secret: {}", hex::encode(&bob_shared[..16]));
    println!("Shared secrets match: {}", alice_shared == bob_shared);

    // ECDH Demo
    println!("\n--- Elliptic Curve Diffie-Hellman ---");
    let alice_ecdh = manager.ecdh_generate().unwrap();
    let bob_ecdh = manager.ecdh_generate().unwrap();

    let alice_ecdh_shared = manager.ecdh_compute_secret(&alice_ecdh.private_key, &bob_ecdh.public_key);
    let bob_ecdh_shared = manager.ecdh_compute_secret(&bob_ecdh.private_key, &alice_ecdh.public_key);

    println!("Alice's ECDH shared secret: {}", hex::encode(&alice_ecdh_shared[..16]));
    println!("Bob's ECDH shared secret: {}", hex::encode(&bob_ecdh_shared[..16]));
    println!("ECDH shared secrets match: {}", alice_ecdh_shared == bob_ecdh_shared);

    // Session Key Derivation
    println!("\n--- Session Key Derivation ---");
    let session_key = manager.derive_session_key(&alice_ecdh_shared, b"session-context", 32);
    println!("Derived session key: {}", hex::encode(&session_key));

    // Authenticated Key Exchange
    println!("\n--- Authenticated Key Exchange ---");
    let alice_public_key = manager.authenticated_key_exchange("[email protected]").unwrap();
    let alice_auth_valid = manager.verify_authenticated_exchange(
        "[email protected]",
        &alice_public_key.0,
        &alice_public_key.1,
    );
    println!("Alice's authenticated key exchange valid: {}", alice_auth_valid);

    // Perfect Forward Secrecy
    println!("\n--- Perfect Forward Secrecy ---");
    let pfs_result = manager.perfect_forward_secrecy_exchange().unwrap();
    println!("PFS Exchange ID: {}", pfs_result.exchange_id);
    println!("PFS Session Key: {}", hex::encode(&pfs_result.session_key));

    // Group Key Exchange
    println!("\n--- Group Key Exchange ---");
    let participants = vec![
        "alice".to_string(),
        "bob".to_string(),
        "charlie".to_string(),
        "diana".to_string(),
    ];

    let group_keys = manager.group_key_exchange(&participants).unwrap();
    println!("Group keys for {} participants:", participants.len());
    for (participant, key) in &group_keys {
        println!("  {}: {}", participant, hex::encode(&key[..16]));
    }

    // Key Rotation Demo
    println!("\n--- Key Rotation ---");
    let original_key = b"original-secret-key-123456";
    let rotated_key1 = manager.rotate_key(original_key, "rotation-1");
    let rotated_key2 = manager.rotate_key(&rotated_key1, "rotation-2");

    println!("Original key: {}", hex::encode(original_key));
    println!("After rotation 1: {}", hex::encode(&rotated_key1[..16]));
    println!("After rotation 2: {}", hex::encode(&rotated_key2[..16]));

    // Post-Quantum Key Exchange
    println!("\n--- Post-Quantum Key Exchange ---");
    let (pq_public, pq_private) = manager.post_quantum_key_exchange().unwrap();
    println!("Post-quantum public key size: {} bytes", pq_public.len());
    println!("Post-quantum private key size: {} bytes", pq_private.len());

    // Run benchmarks
    benchmark_key_exchanges();

    println!("\n=== Key Exchange Demo Complete ===");
}

// Add this to Cargo.toml for hex encoding
// hex = "0.4"