Merge pull request #22 from Matej4545/mg/improved-client

Mg/improved client

.gitignore

@@ -60,4 +60,5 @@ venv/
 
 # Custom data
 server.log
-data/
+data/
+.vscode/

client.py

@@ -5,12 +5,15 @@
 from noise.connection import NoiseConnection, Keypair
 from sys import argv
 import constants
+import argparse
 
 
 class Client:
-    def __init__(self):
+    def __init__(self, server, port):
         self.sock = socket.socket()
         self.key_pair = ED25519().generate_keypair()
+        self.ip = server
+        self.port = port
 
     def exchange_public_keys(self) -> bytes:
         self.sock.send(self.key_pair.public_bytes)
@@ -41,22 +44,92 @@ def send_encrypted_msg(self, message: str):
         encrypted_message = self.noise.encrypt(bytes(message, encoding="UTF-8"))
         self.sock.send(encrypted_message)
 
+    def send_messages(self):
+        print("You can write messages now, [q, quit, exit] to quit:")
+        while True:
+            user_input = input("Message: ")
+            if len(user_input) == 0:
+                print("Please input valid message, [q, quit, exit] to quit.")
+                continue
+            if user_input.lower() in ["q", "quit", "exit"]:
+                return
+            self.communicate(user_input)
+
+    def communicate(self, message):
+        self.sock = socket.socket()
+        self.sock.connect((self.ip, self.port))
+        self.set_connection_keys()
+        self.noise_handshake()
+        self.send_encrypted_msg(message)
+        self.receive_and_decrypt_msg()
+        self.sock.close()
+
     def receive_and_decrypt_msg(self):
         ciphertext = self.sock.recv(constants.CLIENT_PORT)
         plaintext = self.noise.decrypt(ciphertext)
         print(plaintext)
 
+    def register(self):
+        # TODO: read TPMs PCR values and somehow send them to a server
+        print("Registration complete, your identifier is <TOP_SECRET_PCR_HASH>")
+        return True
+
     def run(self, message):
-        self.sock.connect(("localhost", constants.SERVER_PORT))
-        self.set_connection_keys()
-        self.noise_handshake()
-        self.send_encrypted_msg(message)
-        self.receive_and_decrypt_msg()
+        if message:  # One time
+            self.communicate(message)
+        else:  # Multiple messages
+            self.send_messages()
 
 
 if __name__ == "__main__":
-    client = Client()
-    if len(argv) > 1:
-        client.run(argv[1])
-    else:
-        print("Please write message as argument")
+    parser = argparse.ArgumentParser(
+        description="PV204 NoisyTPM - this is a part of team project for PV204. \
+                                                Client app can communicate with server using Noise framework \
+                                                and authenticate via TPM. Please see \
+                                                'https://github.com/Matej4545/PV204-NoisyTPM/' for more info."
+    )
+    parser.add_argument(
+        "-s",
+        "--server",
+        dest="server",
+        metavar="IP",
+        type=str,
+        default="localhost",
+        help="An IP address or hostname of the server.",
+    )
+    parser.add_argument(
+        "-p",
+        "--port",
+        dest="port",
+        metavar="PORT",
+        type=int,
+        default=5555,
+        help="A port where the server is listening.",
+    )
+    parser.add_argument(
+        "-m",
+        "--message",
+        metavar="MESSAGE",
+        dest="message",
+        type=str,
+        nargs="+",
+        help="Specify message as argument. For interactive session please omit.",
+    )
+    parser.add_argument(
+        "-r --register",
+        dest="register",
+        action="store_true",
+        default=False,
+        help="If you are not authenticated or running the app first time, you will need to register.",
+    )
+    args = parser.parse_args()
+
+    try:
+        message = "" if args.message is None else "".join(args.message).strip()
+        client = Client(args.server.strip(), args.port)
+        if args.register:
+            client.register()
+        client.run(message)
+    except Exception as e:
+        print("An error occured! Quitting app.")
+        print(e)