Allow 'startsecure' to be issued at any time

src/main/java/dslab/client/MessageClient.java

@@ -12,6 +12,8 @@ import java.util.Arrays;
 import java.util.Base64;
 import java.util.logging.Logger;
 
+import at.ac.tuwien.dsg.orvell.Shell;
+import at.ac.tuwien.dsg.orvell.annotation.Command;
 import dslab.ComponentFactory;
 import dslab.exception.FailedVerificationException;
 import dslab.util.Config;
@@ -22,6 +24,8 @@ import javax.crypto.spec.SecretKeySpec;
 
 public class MessageClient implements IMessageClient, Runnable {
     private static final Logger logger = Logger.getLogger(MessageClient.class.getName());
+    private final String componentId;
+    private final InputStream consoleIn;
 
     private final String transferHost;
     private final int transferPort;
@@ -38,9 +42,11 @@ public class MessageClient implements IMessageClient, Runnable {
 
     private PrintWriter dmtpOut;
     private BufferedReader dmtpIn;
-    private PrintWriter dmapOut;
+    private PrintStream dmapOut;
     private BufferedReader dmapIn;
 
+    private final Shell shell;
+
     private Cipher aesEncryptCipher;
     private Cipher aesDecryptCipher;
     private byte[] challenge;
@@ -62,6 +68,13 @@ public class MessageClient implements IMessageClient, Runnable {
         this.mailboxUser = config.getString("mailbox.user");
         this.mailboxPassword = config.getString("mailbox.password");
 
+        this.componentId = componentId;
+        this.consoleIn = in;
+
+        this.shell = new Shell(in, out);
+        this.shell.register(this);
+        this.shell.setPrompt(this.componentId + "> ");
+
         logger.info(String.format("TransferHost: %s\nTransferPort: %d\nMailboxHost: %s\nMailboxPort: %d\nTransferEmail: %s\nMailboxUser: %s\nMailboxPassword: %s",
                 transferHost, transferPort, mailboxHost, mailboxPort, transferEmail, mailboxUser, mailboxPassword));
     }
@@ -72,7 +85,7 @@ public class MessageClient implements IMessageClient, Runnable {
             logger.info("Starting connection to MailboxHost on " + this.mailboxHost + " on port " + this.mailboxPort);
             this.dmapSocket = new Socket(this.mailboxHost, this.mailboxPort);
             this.dmapIn = new BufferedReader(new InputStreamReader(this.dmapSocket.getInputStream()));
-            this.dmapOut = new PrintWriter(this.dmapSocket.getOutputStream(), true);
+            this.dmapOut = new PrintStream(this.dmapSocket.getOutputStream(), true);
 
             String input = null;
             String message = null;
@@ -84,10 +97,8 @@ public class MessageClient implements IMessageClient, Runnable {
 
             startSecure();
             loginUser();
+            shell.run();
 
-            while (!Thread.currentThread().isInterrupted() && ((input = this.dmapIn.readLine()) != null)) {
-                // TODO implement commands
-            }
         } catch (IOException e) {
             logger.severe("Could not connect to MailboxHost " + this.mailboxHost + " on port " + this.mailboxPort);
             shutdown();
@@ -103,7 +114,7 @@ public class MessageClient implements IMessageClient, Runnable {
         message = "login " + this.mailboxUser + " " + this.mailboxPassword;
         this.dmapOut.println(getAesCiphertext(message));
         String response = getAesPlaintext(this.dmapIn.readLine());
-        if (!response.startsWith("ok DMAP2.0"))
+        if (!response.startsWith("ok"))
             shutdown();
     }
 
@@ -128,7 +139,6 @@ public class MessageClient implements IMessageClient, Runnable {
                 KeyFactory kf = KeyFactory.getInstance("RSA");
                 // Make generator generate public key from X509 spec
                 serverPublicKey = kf.generatePublic(spec);
-                logger.info("Server's Public Key: " + serverPublicKey);
                 String clientChallenge = generateChallengeMessage(serverPublicKey);
                 // Send clientChallenge to server
                 logger.info("Send clientchallenge to Server: " + clientChallenge);
@@ -274,21 +284,25 @@ public class MessageClient implements IMessageClient, Runnable {
         return (Base64.getEncoder().encodeToString(cipherTextChallenge));
     }
 
+    @Command
     @Override
     public void inbox() {
 
     }
 
+    @Command
     @Override
     public void delete(String id) {
 
     }
 
+    @Command
     @Override
     public void verify(String id) {
 
     }
 
+    @Command
     @Override
     public void msg(String to, String subject, String data) {
         try {
@@ -302,6 +316,7 @@ public class MessageClient implements IMessageClient, Runnable {
         }
     }
 
+    @Command
     @Override
     public void shutdown() {
         if (dmtpSocket != null) {

src/main/java/dslab/mailbox/DMAPConnection.java

@@ -13,14 +13,12 @@ import javax.crypto.spec.SecretKeySpec;
 import java.io.*;
 import java.net.Socket;
 import java.net.SocketException;
-import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.security.*;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.PKCS8EncodedKeySpec;
-import java.security.spec.X509EncodedKeySpec;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.LinkedList;
@@ -41,6 +39,8 @@ public class DMAPConnection implements Runnable {
     private Cipher aesEncryptCipher;
     private Cipher aesDecryptCipher;
 
+    private boolean secure = false;
+
     public DMAPConnection(Socket connection, ConcurrentHashMap<Email, LinkedList<Message>> storage, ConcurrentHashMap<String, String> userStorage, String componentId) {
         this.componentId = componentId;
         this.socket = connection;
@@ -56,20 +56,21 @@ public class DMAPConnection implements Runnable {
             this.out = new PrintWriter(socket.getOutputStream(), true);
             this.in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
             out.println("ok DMAP2.0");
-            startSecure();
+
-            login();
+            loginLoop();
 
             String userInput;
             while (!Thread.currentThread().isInterrupted() && (userInput = in.readLine()) != null) {
                 if ("quit".equals(userInput)) {
                     out.println("ok bye");
-                    shutdown();
+                    loginLoop();
                 } else if ("logout".equals(userInput)) {
                     out.println("ok");
-                    shutdown();
+                    currentUser = null;
+                    loginLoop();
                 } else if ("list".equals(userInput)) {
                     listMessages();
-                } else if ("delete".equals(userInput.split("\\s+")[0])) {
+                } else if (userInput.startsWith("delete")) {
                     if (userInput.split("\\s+").length == 2) {
                         try {
                             deleteMessage(userInput.split("\\s+")[1]);
@@ -79,7 +80,7 @@ public class DMAPConnection implements Runnable {
                     } else {
                         out.println("Please supply a message id to delete!");
                     }
-                } else if ("show".equals(userInput.split("\\s+")[0])) {
+                } else if (userInput.startsWith("show")) {
                     if (userInput.split("\\s+").length == 2) {
                         try {
                             showMessage(userInput.split("\\s+")[1]);
@@ -89,6 +90,8 @@ public class DMAPConnection implements Runnable {
                     } else {
                         out.println("Please supply a message id to show!");
                     }
+                } else if (userInput.startsWith("startsecure")) {
+                    startSecure();
                 } else {
                     out.println("error protocol error");
                     shutdown();
@@ -104,6 +107,57 @@ public class DMAPConnection implements Runnable {
         }
     }
 
+    private void loginLoop() {
+        String userInput;
+
+        try {
+            while (!Thread.currentThread().isInterrupted()) {
+                userInput = in.readLine();
+                if (userInput.startsWith("quit")) {
+                    out.println("ok bye");
+                    shutdown();
+                } else if (userInput.startsWith("login")) {
+                    String[] args = userInput.split("\\s+");
+                    if (args.length != 3)
+                        out.println("Please specify a username and password to login!");
+                    if (this.userStorage.containsKey(args[1])) {
+                        // Check if username exists
+                        if (args[2].equals(this.userStorage.get(args[1]))) {
+                            // Check if password matches
+                            for (Email email : this.storage.keySet()) {
+                                if (args[1].equals(email.getUsername())) {
+                                    // Set current user if login successful
+                                    currentUser = email;
+                                    logger.info("User successfully logged in: " + currentUser.toString());
+                                    out.println("ok");
+                                    return;
+                                }
+                            }
+                        } else {
+                            out.println("error wrong password");
+                        }
+                    } else {
+                        out.println("error unknown user");
+                    }
+                } else if (userInput.startsWith("startsecure")) {
+                    startSecure();
+                } else {
+                    out.println("error not logged in");
+                }
+            }
+        } catch (InterruptedIOException ioe) {
+            logger.info("Received interrupt from parent. Shutting down...");
+            shutdown();
+        } catch (SocketException e) {
+            logger.finer("Received interrupt. Exiting " + this.toString());
+            shutdown();
+        } catch (IOException e) {
+            logger.severe("Failed to get IO-Stream");
+            e.printStackTrace();
+            shutdown();
+        }
+    }
+
     /**
      * Handles the startsecure command issued by the MessageClient.
      */
@@ -111,7 +165,6 @@ public class DMAPConnection implements Runnable {
         String userInput;
 
         try {
-            userInput = in.readLine();  // Should always be "startsecure"
             // Send componentId
             out.println("ok " + componentId);
             // Receive client challenge, secret key and iv
@@ -129,6 +182,7 @@ public class DMAPConnection implements Runnable {
                 logger.severe("Received " + userInput + "from MessageClient. Did not match 'ok'");
                 shutdown();
             }
+            this.secure = true;  // Encrypted communication from now on
         } catch (InterruptedIOException ioe) {
             logger.info("Received interrupt from parent. Shutting down...");
             shutdown();
@@ -230,11 +284,11 @@ public class DMAPConnection implements Runnable {
 
         try {
             while (!Thread.currentThread().isInterrupted()) {
-                userInput = getAesPlaintext(in.readLine());
+                userInput = in.readLine();
-                if ("quit".equals(userInput.split("\\s+")[0])) {
+                if (userInput.startsWith("quit")) {
-                    out.println(getAesCiphertext("ok bye"));
+                    out.println("ok bye");
                     shutdown();
-                } else if ("login".equals(userInput.split("\\s+")[0])) {
+                } else if (userInput.startsWith("login")) {
                     String[] args = userInput.split("\\s+");
                     if (args.length != 3)
                         out.println("Please specify a username and password to login!");
@@ -297,12 +351,14 @@ public class DMAPConnection implements Runnable {
     }
 
     private void listMessages() {
+        logger.info("'list' command received");
         if (storage.get(currentUser).isEmpty()) {
             out.println("You do not have any messages at the moment!");
             return;
         }
 
         for (Message m : storage.get(currentUser)) {
+            logger.info("Printing message from user: " + m.listMessage());
             out.println(m.listMessage());
         }
     }