from cryptography.hazmat.primitives.asymmetric import x25519, ed25519
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import serialization

import os

# -------------------------------------------------
# 1. Key Pair Generation
# -------------------------------------------------

print("\n# 1. Key Pair Generation")
print("# -------------------------------------------------")

# Client ephemeral key (X25519) and signing key (Ed25519)
client_private_key_x25519 = x25519.X25519PrivateKey.generate()
client_public_key_x25519 = client_private_key_x25519.public_key()

client_signing_key = ed25519.Ed25519PrivateKey.generate()
client_verifying_key = client_signing_key.public_key()

# Server ephemeral key (X25519) and signing key (Ed25519)
server_private_key_x25519 = x25519.X25519PrivateKey.generate()
server_public_key_x25519 = server_private_key_x25519.public_key()

server_signing_key = ed25519.Ed25519PrivateKey.generate()
server_verifying_key = server_signing_key.public_key()

print("[+] Key pairs generated successfully.")

# -------------------------------------------------
# 2. Public Key Exchange
# -------------------------------------------------

print("\n# 2. Public Key Exchange")
print("# -------------------------------------------------")

# 2.1 Client sends public key and signature
client_public_key_bytes = client_public_key_x25519.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)
client_signature = client_signing_key.sign(client_public_key_bytes)

print("[2.1] Client sends to media:")
print("    Client Public Key (hex):", client_public_key_bytes.hex())
print("    Client Signature (hex):", client_signature.hex())

# --- Simulate transmission over insecure media ---

# 2.2 Server receives and verifies
print("[2.2] Server receives and verifies Client's key...")

try:
    server_verifying_key_of_client = client_verifying_key  # Trusted out-of-band
    server_verifying_key_of_client.verify(client_signature, client_public_key_bytes)
    print("[+] Server verified Client's public key signature.")
except Exception as e:
    print("[-] Verification failed:", str(e))
    exit(1)

# 2.3 Server sends public key and signature
server_public_key_bytes = server_public_key_x25519.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)
server_signature = server_signing_key.sign(server_public_key_bytes)

print("[2.3] Server sends to media:")
print("    Server Public Key (hex):", server_public_key_bytes.hex())
print("    Server Signature (hex):", server_signature.hex())

# --- Simulate transmission over insecure media ---

# 2.4 Client receives and verifies
print("[2.4] Client receives and verifies Server's key...")

try:
    client_verifying_key_of_server = server_verifying_key  # Trusted out-of-band
    client_verifying_key_of_server.verify(server_signature, server_public_key_bytes)
    print("[+] Client verified Server's public key signature.")
except Exception as e:
    print("[-] Verification failed:", str(e))
    exit(1)

# -------------------------------------------------
# 3. Shared Secret Derivation
# -------------------------------------------------

print("\n# 3. Shared Secret Derivation")
print("# -------------------------------------------------")

# Each side derives the shared secret
server_pub_key_loaded = x25519.X25519PublicKey.from_public_bytes(server_public_key_bytes)
client_shared_secret = client_private_key_x25519.exchange(server_pub_key_loaded)

client_pub_key_loaded = x25519.X25519PublicKey.from_public_bytes(client_public_key_bytes)
server_shared_secret = server_private_key_x25519.exchange(client_pub_key_loaded)

assert client_shared_secret == server_shared_secret, "Shared secret mismatch!"
print("[+] Shared secret established successfully.")
print("Shared secret (hex): ", client_shared_secret.hex())
# -------------------------------------------------
# 4. Key Derivation (HKDF)
# -------------------------------------------------

print("\n# 4. Key Derivation (HKDF)")
print("# -------------------------------------------------")

session_key = HKDF(
    algorithm=hashes.SHA256(),
    length=32,  # 256 bits key for AES-256
    salt=None,
    info=b"handshake data",
).derive(client_shared_secret)

print("[+] Session key derived.")
print("    Session Key (hex):", session_key.hex())

# -------------------------------------------------
# 5. Secure Communication
# -------------------------------------------------

print("\n# 5. Secure Communication")
print("# -------------------------------------------------")

# Encrypt a message
aesgcm = AESGCM(session_key)
nonce = os.urandom(12)  # 96-bit nonce
plaintext = b"Hello from Client to Server!"

ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data=None)

print("[5.1] Client sends encrypted message:")
print("    Nonce (hex):", nonce.hex())
print("    Ciphertext (hex):", ciphertext.hex())

# Decrypt the message
aesgcm_server = AESGCM(session_key)
decrypted_plaintext = aesgcm_server.decrypt(nonce, ciphertext, associated_data=None)

print("[5.2] Server decrypts the message:")
print("    Decrypted Plaintext:", decrypted_plaintext.decode())

print("\n[+] Secure transmission and decryption completed.")

import 'dart:convert';
import 'dart:typed_data';
import 'dart:math';
import 'package:cryptography/cryptography.dart';

void main() async {
  // -------------------------------------------------
  // 1. Key Pair Generation
  // -------------------------------------------------
print('\n# -------------------------------------------------');
  print('# 1. Key Pair Generation');
  print('# -------------------------------------------------');

  final x25519 = X25519();
  final ed25519 = Ed25519();

  // Client generates ephemeral X25519 and signing Ed25519 key pairs
  final clientEphemeralKeyPair = await x25519.newKeyPair();
  final clientSigningKeyPair = await ed25519.newKeyPair();

  // Server generates ephemeral X25519 and signing Ed25519 key pairs
  final serverEphemeralKeyPair = await x25519.newKeyPair();
  final serverSigningKeyPair = await ed25519.newKeyPair();

  print('[+] Key pairs generated successfully.');

  // -------------------------------------------------
  // 2. Public Key Exchange
  // -------------------------------------------------
print('\n# -------------------------------------------------');
  print('# 2. Public Key Exchange');
  print('# -------------------------------------------------');

  // 2.1 Client sends public key and signature
  final clientPublicKeyBytes = await clientEphemeralKeyPair.extractPublicKey().then((k) => k.bytes);
  final clientSignature = await ed25519.sign(
    clientPublicKeyBytes,
    keyPair: clientSigningKeyPair,
  );

  print('[2.1] Client sends to media:');
  print('\tClient Public Key (hex): ${_toHex(clientPublicKeyBytes)}');
  print('\tClient Signature (hex): ${_toHex(clientSignature.bytes)}');

  // --- Simulate transmission ---

  // 2.2 Server receives and verifies
  print('[2.2] Server receives and verifies Client\'s key...');

  final clientVerifyingKey = await clientSigningKeyPair.extractPublicKey();
  try {
    await ed25519.verify(
      clientPublicKeyBytes,
      signature: Signature(clientSignature.bytes, publicKey: clientVerifyingKey),
    );
    print('[+] Server verified Client\'s public key signature.');
  } catch (e) {
    print('[-] Verification failed: $e');
    return;
  }

  // 2.3 Server sends public key and signature
  final serverPublicKeyBytes = await serverEphemeralKeyPair.extractPublicKey().then((k) => k.bytes);
  final serverSignature = await ed25519.sign(
    serverPublicKeyBytes,
    keyPair: serverSigningKeyPair,
  );

  print('[2.3] Server sends to media:');
  print('\tServer Public Key (hex): ${_toHex(serverPublicKeyBytes)}');
  print('\tServer Signature (hex): ${_toHex(serverSignature.bytes)}');

  // --- Simulate transmission ---

  // 2.4 Client receives and verifies
  print('[2.4] Client receives and verifies Server\'s key...');

  final serverVerifyingKey = await serverSigningKeyPair.extractPublicKey();
  try {
    await ed25519.verify(
      serverPublicKeyBytes,
      signature: Signature(serverSignature.bytes, publicKey: serverVerifyingKey),
    );
    print('[+] Client verified Server\'s public key signature.');
  } catch (e) {
    print('[-] Verification failed: $e');
    return;
  }

// -------------------------------------------------
// 3. Shared Secret Derivation
// -------------------------------------------------
print('\n# -------------------------------------------------');
print('# 3. Shared Secret Derivation');
print('# -------------------------------------------------');

// Correct shared secret extraction
final clientSharedSecretKey = await x25519.sharedSecretKey(
  keyPair: clientEphemeralKeyPair,
  remotePublicKey: SimplePublicKey(serverPublicKeyBytes, type: KeyPairType.x25519),
);
final serverSharedSecretKey = await x25519.sharedSecretKey(
  keyPair: serverEphemeralKeyPair,
  remotePublicKey: SimplePublicKey(clientPublicKeyBytes, type: KeyPairType.x25519),
);

// Extract bytes
final clientSharedSecretBytes = await clientSharedSecretKey.extractBytes();
final serverSharedSecretBytes = await serverSharedSecretKey.extractBytes();

// Check
assert(_toHex(clientSharedSecretBytes) == _toHex(serverSharedSecretBytes), 'Shared secret mismatch!');
print('[+] Shared secret established successfully.');
print('\tShared Secret (hex): ${_toHex(clientSharedSecretBytes)}');

// -------------------------------------------------
// 4. Key Derivation (HKDF)
// -------------------------------------------------
print('\n# -------------------------------------------------');
print('# 4. Key Derivation (HKDF)');
print('# -------------------------------------------------');

final hkdf = Hkdf(
  hmac: Hmac.sha256(),
  outputLength: 32,
);

// Derive HKDF nonce securely:
final nonceInput = <int>[];
nonceInput.addAll(clientPublicKeyBytes);
nonceInput.addAll(serverPublicKeyBytes);

final hash = await Sha256().hash(nonceInput);
final hkdfNonce = hash.bytes.sublist(0, 12); // Take first 12 bytes

final sessionKey = await hkdf.deriveKey(
  secretKey: SecretKey(clientSharedSecretBytes),
  nonce: hkdfNonce,
  info: utf8.encode('handshake data'), // constant agreed string
);

final sessionKeyBytes = await sessionKey.extractBytes();

print('[+] Session key derived.');
print('\tSession Key (hex): ${_toHex(sessionKeyBytes)}');

  // -------------------------------------------------
  // 5. Secure Communication
  // -------------------------------------------------
  print('\n# -------------------------------------------------');
  print('# 5. Secure Communication');
  print('# -------------------------------------------------');

  final aesGcm = AesGcm.with256bits();
  final nonce = _randomBytes(12);
  final plaintext = utf8.encode('Hello from Client to Server!');

  final ciphertext = await aesGcm.encrypt(
    plaintext,
    secretKey: sessionKey,
    nonce: nonce,
  );

  print('[5.1] Client sends encrypted message:');
  print('\tNonce (hex): ${_toHex(nonce)}');
  print('\tCiphertext (hex): ${_toHex(ciphertext.cipherText)}');

  final decryptedPlaintext = await aesGcm.decrypt(
    SecretBox(ciphertext.cipherText, nonce: nonce, mac: ciphertext.mac),
    secretKey: sessionKey,
  );

  print('[5.2] Server decrypts the message:');
  print('\tDecrypted Plaintext: ${utf8.decode(decryptedPlaintext)}');

  print('\n[+] Secure transmission and decryption completed.');
}

String _toHex(List<int> bytes) {
  return bytes.map((b) => b.toRadixString(16).padLeft(2, '0')).join();
}

List<int> _randomBytes(int length) {
  final rand = Random.secure();
  return List<int>.generate(length, (_) => rand.nextInt(256));
}