package io.github.bindon;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.List;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Hex;

public class ECDHServer {
	public static void testECDHServer() {
		ServerSocket serverSocket = null;
		Socket clientSocket = null;
		BufferedReader reader = null;
		PrintWriter writer = null;
		
		try {
			serverSocket = new ServerSocket(32768);
			clientSocket = serverSocket.accept();
			reader = new BufferedReader(new InputStreamReader(clientSocket.getInputStream()));
			writer = new PrintWriter(new OutputStreamWriter(clientSocket.getOutputStream()));
	
			// Generate ephemeral ECDH keypair
			KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC");
			keyPairGenerator.initialize(256);
			KeyPair serverKeyPair = keyPairGenerator.generateKeyPair();
			byte[] serverPublicKeyBytes = serverKeyPair.getPublic().getEncoded();
			
			// Send server public key
			writer.println(Base64.getUrlEncoder().encodeToString(serverPublicKeyBytes));
			writer.flush();
			
			// Receive client public key
			byte[] clientPublicKeyBytes = Base64.getUrlDecoder().decode(reader.readLine());
			System.out.println("[Server] " + Hex.encodeHexString(clientPublicKeyBytes));
			KeyFactory clientKeyFactory = KeyFactory.getInstance("EC");
			X509EncodedKeySpec clientKeySpec = new X509EncodedKeySpec(clientPublicKeyBytes);
			PublicKey clientPublicKey = clientKeyFactory.generatePublic(clientKeySpec);
			
			// Perform key agreement
			KeyAgreement keyAgreement = KeyAgreement.getInstance("ECDH");
			keyAgreement.init(serverKeyPair.getPrivate());
			keyAgreement.doPhase(clientPublicKey, true);
			
			// Generate shared secret
			byte[] sharedSecret = keyAgreement.generateSecret();
			System.out.println("[Server] Shared secret: " + Hex.encodeHexString(sharedSecret));
			
			// Derive a key from the shared secret and both public keys
			MessageDigest hash = MessageDigest.getInstance("SHA-256");
			hash.update(sharedSecret); // Simple deterministic ordering
			List<ByteBuffer> keys = Arrays.asList(
					ByteBuffer.wrap(serverPublicKeyBytes), ByteBuffer.wrap(clientPublicKeyBytes));
			Collections.sort(keys);
			hash.update(keys.get(0));
			hash.update(keys.get(1));
			
			byte[] derivedKey = hash.digest();
			System.out.println("[Server] Final key: " + Hex.encodeHexString(derivedKey));
			
			// AES-GCM Decrypt and Send data using derivedKey
			byte[] iv = new byte[12];
			Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
			SecretKeySpec keySpec = new SecretKeySpec(derivedKey, "AES");
			GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(128, iv);
			cipher.init(Cipher.DECRYPT_MODE, keySpec, gcmParameterSpec);
			byte[] plaintext = cipher.doFinal(Base64.getUrlDecoder().decode(reader.readLine()));
			System.out.println("[Server] Received data : " + new String(plaintext, StandardCharsets.UTF_8));
		} catch(Exception e) {
			e.printStackTrace();
		} finally {
			if(reader != null) { try { reader.close(); } catch(Exception ignore) {} }
			if(writer != null) { try { writer.close(); } catch(Exception ignore) {} }
			if(clientSocket != null) { try { clientSocket.close(); } catch(Exception ignore) {} }
			if(serverSocket != null) { try { serverSocket .close(); } catch(Exception ignore) {} }
		}
	}
}

package io.github.bindon;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.List;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Hex;

public class ECDHClient {
	public static void testECDHClient() {
		Socket socket = null;
		PrintWriter writer = null;
		BufferedReader reader = null;
		
		try {
			socket = new Socket("127.0.0.1", 32768);
			reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
			writer = new PrintWriter(new OutputStreamWriter(socket.getOutputStream()));
			
			// Generate ephemeral ECDH keypair
			KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC");
			keyPairGenerator.initialize(256);
			KeyPair clientKeyPair = keyPairGenerator.generateKeyPair();
			byte[] clientPublicKeyBytes = clientKeyPair.getPublic().getEncoded();
			
			// Receive server public key
			byte[] serverPublicKeyBytes = Base64.getUrlDecoder().decode(reader.readLine());
			System.out.println("[Client] " + Hex.encodeHexString(serverPublicKeyBytes));
			KeyFactory serverKeyFactory = KeyFactory.getInstance("EC");
			X509EncodedKeySpec serverKeySpec = new X509EncodedKeySpec(serverPublicKeyBytes);
			PublicKey serverPublicKey = serverKeyFactory.generatePublic(serverKeySpec);

			// Send client public key
			writer.println(Base64.getUrlEncoder().encodeToString(clientPublicKeyBytes));
			writer.flush();
	
			// Perform key agreement
			KeyAgreement keyAgreement = KeyAgreement.getInstance("ECDH");
			keyAgreement.init(clientKeyPair.getPrivate());
			keyAgreement.doPhase(serverPublicKey, true);
			
			// Read shared secret
			byte[] sharedSecret = keyAgreement.generateSecret();
			System.out.println("[Client] Shared secret: " + Hex.encodeHexString(sharedSecret));
			
			// Derive a key from the shared secret and both public keys
			MessageDigest hash = MessageDigest.getInstance("SHA-256");
			hash.update(sharedSecret); // Simple deterministic ordering
			List<ByteBuffer> keys = Arrays.asList(
					ByteBuffer.wrap(clientPublicKeyBytes), ByteBuffer.wrap(serverPublicKeyBytes));
			Collections.sort(keys);
			hash.update(keys.get(0));
			hash.update(keys.get(1));
			
			byte[] derivedKey = hash.digest();
			System.out.println("[Client] Final key: " + Hex.encodeHexString(derivedKey));
			
			// AES-GCM Encrypt and Send data using derivedKey
			byte[] iv = new byte[12];
			Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
			SecretKeySpec keySpec = new SecretKeySpec(derivedKey, "AES");
			GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(128, iv);
			cipher.init(Cipher.ENCRYPT_MODE, keySpec, gcmParameterSpec);
			byte[] ciphertext = cipher.doFinal("bindon".getBytes(StandardCharsets.UTF_8));
			writer.println(Base64.getUrlEncoder().encodeToString(ciphertext));
			writer.flush();
		} catch(Exception e) {
			e.printStackTrace();
		} finally {
			if(reader != null) { try { reader.close(); } catch(Exception ignore) {} }
			if(writer != null) { try { writer.close(); } catch(Exception ignore) {} }
			if(socket != null) { try { socket.close(); } catch(Exception ignore) {} }
		}
	}
}