Step 1 completed

src/main/java/step1/Main.java

@@ -3,6 +3,7 @@ package step1;
 import javax.crypto.Cipher;
 import javax.crypto.KeyGenerator;
 import javax.crypto.SecretKey;
+import javax.crypto.spec.GCMParameterSpec;
 import javax.crypto.spec.IvParameterSpec;
 import javax.crypto.spec.SecretKeySpec;
 import java.io.*;
@@ -23,7 +24,6 @@ public class Main {
 
 
         try (Socket socket = new Socket()) {
-
             socket.connect(new InetSocketAddress("127.0.0.1", 8888));
 
             //Saving the buffer reference to read it later
@@ -33,40 +33,78 @@ public class Main {
             //Waiting for welcome message so it does
             //not pollute the fucking stream
             String welcome = readResponse(reader);
-            System.out.println("Welcome message: " + welcome);
+            displayReceived("Welcome message: " + welcome);
 
             //Creating the key
             SecretKey key = get3DESKey();
             byte[] encodedKey = key.getEncoded();
             String base64Key = Base64.getEncoder().encodeToString(encodedKey);
+            displayInfo("Generated key: " + base64Key);
 
             //Sending the key
             send(socket, base64Key);
-            System.out.println("Key sent to server.");
+            displaySent("Key sent to server : " + base64Key);
 
             //Reading the message sent by the server
             String text = readResponse(reader);
-            System.out.println("Received message: " + text);
+            displayReceived("Received message: " + text);
 
             //Encrypting the message
             byte[] encryptedMessage = encrypt(key, text);
             String base64Message = Base64.getEncoder().encodeToString(encryptedMessage);
-            System.out.println("Encrypted message to send : " + base64Message);
+            displayInfo("Encrypted message : " + base64Message);
 
             //Sending encrypted message
             send(socket, base64Message);
-            System.out.println("Encrypted message sent to server.");
+            displaySent("Encrypted message sent to server : " + base64Message);
 
-            //Reading the parameters sent by the server
+            //Reading the base 64 key sent by the server
-            String parameters = readResponse(reader);
+            String receivedKey = readResponse(reader);
-            System.out.println("Received parameters: " + parameters);
+            displayReceived("Received base64 key: " + receivedKey);
 
+            //Reading the base 64 IV sent by the server
+            String base64IV = readResponse(reader);
+            displayReceived("Received base 64 IV: " + base64IV);
+
+            //Reading the AAD sent by the server
+            String AAD = readResponse(reader);
+            displayReceived("Received AAD: " + AAD);
+
+            //Reading the base 64 encrypted message sent by the server
+            String base64ReceivedMessage = readResponse(reader);
+            displayReceived("Received base 64 encrypted message: " + base64ReceivedMessage);
+
+            String decryptedMessage = decrypt(receivedKey, base64IV, AAD, base64ReceivedMessage);
+            displayInfo("Decrypted message : " + decryptedMessage);
+
+            send(socket, decryptedMessage);
+            displaySent("Decrypted message sent to server : " + decryptedMessage);
+
+            displayInfo("Process completed. Check out the server output to know if it was successful.");
 
         } catch (Exception e) {
             e.printStackTrace();
         }
     }
 
+    public static void displayInfo(String message){
+        final String BLUE = "\u001B[34m";
+        final String RESET = "\u001B[0m";
+        System.out.println(BLUE + "Info >>> " + message + RESET);
+    }
+
+    public static void displayReceived(String message){
+        String GREEN = "\u001B[32m";
+        String RESET = "\u001B[0m";
+        System.out.println(GREEN + "Received >>> " + message + RESET);
+    }
+
+    public static void displaySent(String message){
+        final String RED = "\u001B[31m";
+        final String RESET = "\u001B[0m";
+        System.out.println(RED + "Sent >>> " + message + RESET);
+    }
+
     public static void send(Socket socket, String data) throws IOException {
         OutputStream output = socket.getOutputStream();
         output.write(String.format("%s\r\n", data).getBytes(StandardCharsets.UTF_8));
@@ -91,24 +129,43 @@ public class Main {
 
     }
 
-    public static String decrypt(byte[] keyBytes, String base64IvAndCiphertext) throws Exception {
+    public static String decrypt(String base64Key, String base64IV, String aad, String base64Ciphertext) throws Exception {
-        SecretKey key = new SecretKeySpec(keyBytes, ALGORITHM);
 
-        byte[] ivPlusCipher = Base64.getDecoder().decode(base64IvAndCiphertext);
+        final int GCM_TAG_LENGTH = 128;
-        if (ivPlusCipher.length < 8) throw new IllegalArgumentException("Data too short");
 
-        byte[] iv = new byte[8];
+        // Decode Base64 inputs
-        System.arraycopy(ivPlusCipher, 0, iv, 0, 8);
+        byte[] key = Base64.getDecoder().decode(base64Key);
-        IvParameterSpec ivSpec = new IvParameterSpec(iv);
+        byte[] iv = Base64.getDecoder().decode(base64IV);
+        byte[] cipherMessage = Base64.getDecoder().decode(base64Ciphertext);
 
-        byte[] ciphertext = new byte[ivPlusCipher.length - 8];
+        // Split ciphertext and authentication tag (last 16 bytes)
-        System.arraycopy(ivPlusCipher, 8, ciphertext, 0, ciphertext.length);
+        int tagLength = GCM_TAG_LENGTH / 8;
+        int ciphertextLength = cipherMessage.length - tagLength;
+        byte[] ciphertext = new byte[ciphertextLength];
+        byte[] tag = new byte[tagLength];
 
-        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
+        System.arraycopy(cipherMessage, 0, ciphertext, 0, ciphertextLength);
-        cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
+        System.arraycopy(cipherMessage, ciphertextLength, tag, 0, tagLength);
 
-        byte[] plaintextBytes = cipher.doFinal(ciphertext);
+        // Combine ciphertext + tag
-        return new String(plaintextBytes, StandardCharsets.UTF_8);
+        byte[] ciphertextWithTag = new byte[cipherMessage.length];
+        System.arraycopy(cipherMessage, 0, ciphertextWithTag, 0, cipherMessage.length);
+
+        // Create AES key and GCM spec
+        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
+        GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
+
+        // Initialize cipher
+        Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
+        cipher.init(Cipher.DECRYPT_MODE, keySpec, gcmSpec);
+
+        // Add AAD (if any)
+        if (aad != null && !aad.isEmpty()) {
+            cipher.updateAAD(aad.getBytes());
         }
 
+        // Decrypt and return plaintext
+        byte[] plaintext = cipher.doFinal(ciphertextWithTag);
+        return new String(plaintext, "UTF-8");
+    }
 }