From b29e80a4af154e7fe7d96d1ef3370ce28c2acba4 Mon Sep 17 00:00:00 2001
From: WickedJack99 <wickedjack1@proton.me>
Date: Fri, 1 Dec 2023 22:39:55 +0100
Subject: [PATCH] Moved information into external classes, added certificate
 validation, implemented receive function.

---
 src/main/java/connect/src/TLSClient.java | 69 +++++++++++++++++++-----
 1 file changed, 57 insertions(+), 12 deletions(-)

diff --git a/src/main/java/connect/src/TLSClient.java b/src/main/java/connect/src/TLSClient.java
index 744c40f..66469a1 100644
--- a/src/main/java/connect/src/TLSClient.java
+++ b/src/main/java/connect/src/TLSClient.java
@@ -1,38 +1,55 @@
 package connect.src;
 
 import java.io.BufferedReader;
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.io.PrintWriter;
 import java.net.UnknownHostException;
+import java.security.KeyManagementException;
 import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+import java.security.cert.CertificateException;
 
+import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLSocket;
 import javax.net.ssl.SSLSocketFactory;
+import javax.net.ssl.TrustManagerFactory;
 
 public final class TLSClient implements Client {
 
-    private ClientInformation clientInformation;
-    // Path to keystore file
-    private String keyStoreFile = "";
-    // Password to access the keystore file
-    private String keyStoreFilePassword = "";
-    // Alias password
-    private String keyStoreAliasPassword = "";
-
-    private KeyStore keyStore;
+    private TLSClientInformation clientInformation;
 
     private SSLSocket sslSocket;
     private BufferedReader reader;
     private PrintWriter writer;
 
+    public TLSClient(ClientInformation clientInformation) {
+        this.setClientData(clientInformation);
+    }
+
     public void setClientData(ClientInformation clientInformation) {
-        this.clientInformation = clientInformation;
+        this.clientInformation = (TLSClientInformation)clientInformation;
     }
 
     public void createConnection() {
         try {
-            SSLSocketFactory sslSocketFactory = (SSLSocketFactory) SSLSocketFactory.getDefault();
+
+            KeyStore truststore = KeyStore.getInstance("JKS");
+            truststore.load(new FileInputStream(this.clientInformation.truststorePath()), this.clientInformation.truststorePassword());
+
+            // Create a TrustManager that trusts the certificates in the truststore
+            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
+            trustManagerFactory.init(truststore);
+
+            // Create an SSL context with the trust manager
+            SSLContext sslContext = SSLContext.getInstance("TLS");
+            sslContext.init(null, trustManagerFactory.getTrustManagers(), null);
+
+            // Set the SSL context on the SSLSocketFactory
+            SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
+            
             sslSocket = (SSLSocket) sslSocketFactory.createSocket("127.0.0.1", 5000);
             reader = new BufferedReader(new InputStreamReader(sslSocket.getInputStream()));
             writer = new PrintWriter(sslSocket.getOutputStream());
@@ -42,13 +59,27 @@ public final class TLSClient implements Client {
         } catch (IOException e) {
             // TODO Auto-generated catch block
             e.printStackTrace();
+        } catch (NoSuchAlgorithmException e) {
+            // TODO Auto-generated catch block
+            e.printStackTrace();
+        } catch (CertificateException e) {
+            // TODO Auto-generated catch block
+            e.printStackTrace();
+        } catch (KeyStoreException e) {
+            // TODO Auto-generated catch block
+            e.printStackTrace();
+        } catch (KeyManagementException e) {
+            // TODO Auto-generated catch block
+            e.printStackTrace();
         }
     }
 
     public void sendData(Data data) {
         if (sslSocket != null && sslSocket.isConnected() && reader != null && writer != null) {
+            writer.println("Hello Server!");
+            writer.flush();
             if (null != data) {
-                writer.println("Hello Server!");
+                writer.println(data.toString());
                 writer.flush();
             }
         } else {
@@ -56,6 +87,20 @@ public final class TLSClient implements Client {
         }
     }
 
+    public void receiveData() {
+        try {
+            if (null != reader) {
+                String response;
+                while ((response = reader.readLine()) != null) {
+                    System.out.println(response);
+                }
+            }
+        } catch (IOException e) {
+            // TODO Auto-generated catch block
+            e.printStackTrace();
+        }
+    }
+
     public void closeConnection() {
         try {
             if (null != reader) {