From 7778a91e34332a3daa818722b4e7a15ca98005d5 Mon Sep 17 00:00:00 2001
From: RIEFOLO MARIANO <MARIANO.RIEFOLO@TRIENNIO.LOCAL>
Date: Tue, 9 Apr 2024 10:39:42 +0200
Subject: [PATCH] Encrypt messages with recipient public key

---
 src/main/java/controllers/Database.java     | 28 +++++++++++++++++++++
 src/main/java/controllers/ServerThread.java |  4 +++
 src/main/java/models/PublicKey.java         | 14 +++++++++++
 src/main/java/views/ClientSendThread.java   | 11 +++++++-
 4 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/models/PublicKey.java

diff --git a/src/main/java/controllers/Database.java b/src/main/java/controllers/Database.java
index c2ab4d8..5d79f36 100644
--- a/src/main/java/controllers/Database.java
+++ b/src/main/java/controllers/Database.java
@@ -13,6 +13,7 @@ import java.sql.*;
 import java.time.LocalDateTime;
 import java.util.ArrayList;
 import java.util.List;
+import models.PublicKey;
 
 public class Database {
     public static boolean registerAccount(String username, String password, BigInteger e, String d, BigInteger n) {
@@ -126,6 +127,33 @@ public class Database {
             return false;
         }
     }
+    
+    public static PublicKey getPublicKey(int accountId) {
+        try (Connection connection = getConnection()) {
+            if (connection == null) return null;
+
+            try (PreparedStatement statement = connection.prepareStatement("""
+            SELECT e, n
+            FROM accounts
+            WHERE id = ?
+        """)) {
+                statement.setInt(1, accountId);
+                ResultSet resultSet = statement.executeQuery();
+                if (resultSet.next()) {
+                    BigInteger e = new BigInteger(resultSet.getString(1));
+                    BigInteger n = new BigInteger(resultSet.getString(2));
+                    return new PublicKey(e, n);
+                }
+                return null;
+            } catch (SQLException e) {
+                System.err.println("Error while trying to get the public key: " + e);
+                return null;
+            }
+        } catch (SQLException e) {
+            System.err.println("Error while trying to open a connection: " + e);
+            throw new RuntimeException(e);
+        }
+    }
 
     public static List<Conversation> getConversations(int sender_id) {
         try (Connection connection = getConnection()) {
diff --git a/src/main/java/controllers/ServerThread.java b/src/main/java/controllers/ServerThread.java
index 9ea6e18..c6297f6 100644
--- a/src/main/java/controllers/ServerThread.java
+++ b/src/main/java/controllers/ServerThread.java
@@ -10,6 +10,7 @@ import java.net.Socket;
 import java.util.Arrays;
 
 import static controllers.Database.*;
+import models.PublicKey;
 
 public class ServerThread extends Thread {
     private final Socket client;
@@ -108,6 +109,9 @@ public class ServerThread extends Thread {
         int recipientId = getIdFromUsername(dUsername);
         Database.addConversation(clientId, recipientId);
 
+        PublicKey pk = Database.getPublicKey(recipientId);
+        sendEncrypted(pk.e().toString());
+        sendEncrypted(pk.n().toString());
 
         for (;;) {
             try {
diff --git a/src/main/java/models/PublicKey.java b/src/main/java/models/PublicKey.java
new file mode 100644
index 0000000..af7c1c8
--- /dev/null
+++ b/src/main/java/models/PublicKey.java
@@ -0,0 +1,14 @@
+/*
+ * Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license
+ * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template
+ */
+package models;
+
+import java.math.BigInteger;
+
+/**
+ *
+ * @author mariano.riefolo
+ */
+public record PublicKey(BigInteger e, BigInteger n) {
+}
diff --git a/src/main/java/views/ClientSendThread.java b/src/main/java/views/ClientSendThread.java
index f035a5e..5417b8d 100644
--- a/src/main/java/views/ClientSendThread.java
+++ b/src/main/java/views/ClientSendThread.java
@@ -62,10 +62,19 @@ public class ClientSendThread extends Thread{
             Logger.getLogger(ClientSendThread.class.getName()).log(Level.SEVERE, null, ex);
         }
         
+        BigInteger recipientE, recipientN;
+        try {
+            recipientE = new BigInteger(rsa.decrypt(in.readLine()));
+            recipientN = new BigInteger(rsa.decrypt(in.readLine()));
+        } catch (IOException ex) {
+            Logger.getLogger(ClientSendThread.class.getName()).log(Level.SEVERE, null, ex);
+            return;
+        }
+        
         while (true) {
             try {
                 s = br.readLine();
-                send(rsa.encrypt(s, se, sn));
+                send(rsa.encrypt(s, recipientE, recipientN));
             } catch (IOException ex) {
                 Logger.getLogger(ClientSendThread.class.getName()).log(Level.SEVERE, null, ex);
             }