Code refactoring

.idea/cmd-chat.iml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/venv" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (cmd-chat)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/cmd-chat.iml" filepath="$PROJECT_DIR$/.idea/cmd-chat.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

client/client.py

@@ -1,125 +1,103 @@
 import os 
 import time
-import rsa 
 import platform
-import requests
 import threading
 from colorama import init
 from colorama import Fore
-from cryptography.fernet import Fernet
 from websocket import create_connection
+from core.crypto import RSAService
+
 
 init()
 
-OS = "Windows"
-if "Linux" in str(platform.platform()):
-    OS = "Linux"
 
-class Client:
+class Client(RSAService):
 
-    def _key_gen(self) -> None:
+    def __init__(self, server: str, port: int, username: str):
-        (pubkey, privkey) = rsa.newkeys(512)
+        super().__init__()
-        with open("private.pem", "wb") as f:
-            f.write(privkey.save_pkcs1())
-        with open("public.pem", "wb") as f:
-            f.write(pubkey.save_pkcs1())
-
-    def __init__(self):
         # Server info 
-        self.server = input("server ip: \n") 
+        self.server = server
-        self.port = input("server port: \n") 
+        self.port = port
-        self.username = input("your username: \n") 
+        self.username = username
         # Urls 
         self.base_url = f"http://{self.server}:{self.port}"
         self.talk_url = f"{self.base_url}/talk"
         self.info_url = f"{self.base_url}/update"
         self.key_url = f"{self.base_url}/get_key"
         self.ws_url = f"ws://{self.server}:{self.port}"
-        # Keys 
+
-        self.pubkey = None 
+    def __get_os(self) -> str:
-        self.privkey = None
+        if "Linux" in str(platform.platform()):
-        self.symetric_key = None 
+            return "Linux"
-        self.fernet = None
+        return "Windows"
 
     def send_info(self):
         ws = create_connection(f"{self.ws_url}/talk")
         while True:
-            user_input = input("You're message: ")
+            try:
-            message = f'{self.username}: {user_input}'
+                user_input = input("You're message: ")
-            socket_message = str({
+                message = f'{self.username}: {user_input}'
-                "text": self.fernet.encrypt(message.encode()),
+                socket_message = str({
-                "username": self.username
+                    "text": self._encrypt(message),
-            })
+                    "username": self.username
-            ws.send(payload=socket_message.encode())
+                })
+                ws.send(payload=socket_message.encode())
+            except KeyboardInterrupt:
+                ws.close()
+                quit()
+            except Exception as exc:
+                ws.close()
+                print("Something went wrong! ", exc)
+                quit()
 
     def print_message(self, message: str) -> str:
         message = message.split(":")
         if message[0] == self.username:
             return Fore.MAGENTA + message[0] + ": " + message[1] + Fore.WHITE
+        return message[0] + ": " + message[1] + Fore.WHITE
+
+    def __clear_console(self):
+        # For windows clear command its cls
+        # For linux clear command its clear
+        if self.__get_os() == "Linux":
+            os.system("clear")
         else:
-            return message[0] + ": " + message[1] + Fore.WHITE 
+            os.system("cls")
 
     def update_info(self):
         ws = create_connection(f"{self.ws_url}/update")
         last_try = None
         while True:
-            time.sleep(0.05)
+            try:
-            r = ws.recv()
+                time.sleep(0.05)
-            if last_try == eval(r):
+                r = eval(ws.recv())
-                continue 
+                if last_try == r:
-            last_try = eval(r)
+                    continue
-            # For windows clear command its cls 
+                last_try = r
-            # For linux clear command its clear
+                self.__clear_console()
-            if OS == "Linux":
+                if len(last_try['status']) > 0:
-                os.system("clear")
+                    for i, msg in enumerate(last_try["status"]):
-            else:
+                        actual_message = self._decrypt(msg)
-                os.system("cls")
+                        if i == 0:
-            if len(last_try['status']) > 0:
+                            for user in last_try["users_in_chat"]:
-                for i, msg in enumerate(last_try["status"]):
+                                print("IP:", Fore.MAGENTA + user.split(",")[0] + Fore.WHITE)
-                    actual_message = self.fernet.decrypt(
+                                print("USERNAME: ", Fore.GREEN + user.split(",")[1] + Fore.WHITE)
-                        msg.encode()
+                            print(f"\n{self.print_message(actual_message)}")
-                    ).decode("utf-8")
+                        else:
-                    if i == 0:
+                            print(f"{self.print_message(actual_message)}")
-                        users = last_try["users_in_chat"]
+            except KeyboardInterrupt:
-                        for user in users:
+                ws.close()
-                            ip = user.split(",")[0]
+                quit()
-                            username = user.split(",")[1]
+            except Exception as exc:
-                            print("IP:", Fore.MAGENTA + ip + Fore.WHITE)
+                ws.close()
-                            print("USERNAME: ", Fore.GREEN + username + Fore.WHITE)
+                print("Something went wrong! ", exc)
-                        print()
+                quit()
-                        print(f"{self.print_message(actual_message)}")
-                    else:
-                        print(f"{self.print_message(actual_message)}")
-
-    def _key_request(self) -> None:
-        with open('private.pem', 'rb') as f:
-            self.privkey = rsa.PrivateKey.load_pkcs1(f.read())
-        with open("public.pem", 'rb') as f:
-            r = requests.get(
-                self.key_url, 
-                data={
-                    "pubkey": f.read(), 
-                    "username": self.username
-                }, 
-                stream=True
-            )
-            message = r.raw.read(999)
-            self.symetric_key = rsa.decrypt(message, self.privkey)
-            self.fernet = Fernet(self.symetric_key)
-
-    def _remove_keys(self) -> None:
-        os.remove("private.pem")
-        os.remove("public.pem")
 
     def _validate_keys(self) -> None:
-        self._key_gen()
+        self._request_key(self.key_url, self.username)
-        self._key_request()
-        with open('public.pem', "rb") as f:
-            first_key = f.read()   
-        self.pubkey = rsa.PublicKey.load_pkcs1(first_key)
         self._remove_keys()
 
-    def __call__(self):
+    def run(self):
         # Running two threads,
         # One for sending info
         # Second one for updating info
@@ -133,4 +111,8 @@ class Client:
 
 
 if __name__ == '__main__':
-    Client()()
+    Client(
+        server=input("server ip:\n"),
+        port=int(input("server port: \n")),
+        username=input("username:\n").replace(" ", "").lower()
+    ).run()

client/core/abs/abs_crypto.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+
+
+class CryptoService(ABC):
+
+    @abstractmethod
+    def _encrypt(self, message: str) -> str:
+        raise NotImplementedError("Need to implement encrypt method")
+
+    @abstractmethod
+    def _decrypt(self, message: str) -> str:
+        raise NotImplementedError("Need to implement decrypt method")
+
+    @abstractmethod
+    def _request_key(self, url: str, username: str):
+        raise NotImplementedError("Need to implement request key method")
+
+    @abstractmethod
+    def _generate_keys(self):
+        raise NotImplementedError("Need to implement generate keys method")
+
+    @abstractmethod
+    def _get_generated_keys(self) -> list[str]:
+        raise NotImplementedError("Need to implement get generated keys method")
+
+    @abstractmethod
+    def _remove_keys(self):
+        raise NotImplementedError("Need to implement remove keys method")

client/core/crypto.py

@@ -0,0 +1,65 @@
+import os
+import rsa
+import requests
+from cryptography.fernet import Fernet
+from core.abs.abs_crypto import CryptoService
+
+
+class RSAService(CryptoService):
+    def __init__(self):
+        self.public_key = None
+        self.private_key = None
+        self.symmetric_key = None
+        self.fernet = None
+        self.private_key_name = "private.pem"
+        self.public_key_name = "public.pem"
+
+        self.keys_path: list[str] = []
+        self._generate_keys()
+
+    def _encrypt(self, message: str) -> str:
+        return self.fernet.encrypt(message.encode())
+
+    def _decrypt(self, message: str) -> str:
+        return self.fernet.decrypt(message.encode()).decode("utf-8")
+
+    def _request_key(self, url: str, username: str):
+        data = {
+            "pubkey": self._open_generated_file(self.public_key_name),
+            "username": username
+        }
+        r = requests.get(url, data=data, stream=True)
+        message = r.raw.read(999)
+        self.symmetric_key = rsa.decrypt(message, self.private_key)
+        self.fernet = Fernet(self.symmetric_key)
+
+    def __update_keys_path(self, path_list: list[str]) -> None:
+        for file in path_list:
+            self.keys_path.append(file)
+
+    def __write_generated_key(self, name: str, key) -> None:
+        with open(name, "wb") as f:
+            f.write(key.save_pkcs1())
+
+    def _open_generated_file(self, name: str) -> bytes:
+        with open(name, "rb") as f:
+            return f.read()
+
+    def _generate_keys(self):
+        (public_key, private_key) = rsa.newkeys(512)
+        self.__write_generated_key(self.private_key_name, private_key)
+        self.__write_generated_key(self.public_key_name, public_key)
+        self.private_key = rsa.PrivateKey.load_pkcs1(
+            self._open_generated_file(self.private_key_name)
+        )
+        self.public_key = rsa.PublicKey.load_pkcs1(
+            self._open_generated_file(self.public_key_name)
+        )
+        self.__update_keys_path(["public.pem", "private.pem"])
+
+    def _get_generated_keys(self):
+        return self.private_key, self.public_key
+
+    def _remove_keys(self):
+        for key in self.keys_path:
+            os.remove(key)

private.pem

@@ -0,0 +1,9 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIBPAIBAAJBAJYbSpDoRC8IBswfwPT2hu8m6w9lOAXCY3Oq8srtsvLhrknrdXsb
+FWeWwikTJGkuWff5zNZXC8LN72+OdYXJGEkCAwEAAQJACUYcYEGJXOKBEQFxOXE1
+uvbLlQLq6CgvXskUAQeYeQXkU+yAuFo72ZqIP4TDs+xn1RcjzLjRCBshsS+Yru3d
+xQIjAK1PezX4YtRWdP2PkTpNtTpbIFgua88MKFYlW+jro4pDonMCHwDduaOcFhZL
+Gszx3inSzEal9NvBIrfSwqwMGVSyn1MCIgE9gV84gNSOLdYmsd5d8f8R6eBXrLPV
+nXBIYij/jrMNmtsCHwCLb3kCaalvZdVIrYvjsu8i4o9oL+smMaJ8oVlwU10CImFC
+uFlA2nG0Zgocwy6WA7laZKd6clskEdTxRwoOQRdpfSM=
+-----END RSA PRIVATE KEY-----

public.pem

@@ -0,0 +1,4 @@
+-----BEGIN RSA PUBLIC KEY-----
+MEgCQQCWG0qQ6EQvCAbMH8D09obvJusPZTgFwmNzqvLK7bLy4a5J63V7GxVnlsIp
+EyRpLln3+czWVwvCze9vjnWFyRhJAgMBAAE=
+-----END RSA PUBLIC KEY-----