Code refactoring
This commit is contained in:
mirai
@@ -0,0 +1,159 @@
|
||||
import os
|
||||
import time
|
||||
import platform
|
||||
import threading
|
||||
|
||||
from colorama import init
|
||||
from websocket import create_connection
|
||||
|
||||
from cmd_chat.client.core.crypto import RSAService
|
||||
from cmd_chat.client.config import (
|
||||
COLORS
|
||||
)
|
||||
|
||||
|
||||
init()
|
||||
|
||||
|
||||
class Client(RSAService):
|
||||
|
||||
def __init__(self, server: str, port: int, username: str):
|
||||
super().__init__()
|
||||
# Server info
|
||||
self.server = server
|
||||
self.port = port
|
||||
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}"
|
||||
self.close_response = str({
|
||||
"action": "close",
|
||||
"username": self.username
|
||||
})
|
||||
|
||||
def __get_os(self) -> str:
|
||||
""" checking what kind of platform you need
|
||||
"""
|
||||
if "Linux" in str(platform.platform()):
|
||||
return "Linux"
|
||||
return "Windows"
|
||||
|
||||
def send_info(self):
|
||||
""" sending message to websocket
|
||||
"""
|
||||
ws = create_connection(f"{self.ws_url}/talk")
|
||||
while True:
|
||||
try:
|
||||
user_input = input("You're message: ")
|
||||
message = f'{self.username}: {user_input}'
|
||||
socket_message = str({
|
||||
"text": self._encrypt(message),
|
||||
"username": self.username
|
||||
})
|
||||
ws.send(
|
||||
payload=socket_message.encode()
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
ws.send(self.close_response)
|
||||
ws.close()
|
||||
quit()
|
||||
except Exception as exc:
|
||||
ws.send(self.close_response)
|
||||
ws.close()
|
||||
print("Something went wrong! ", exc)
|
||||
quit()
|
||||
|
||||
def __print_message(self, message: str) -> str:
|
||||
""" generating string with message in required format
|
||||
"""
|
||||
message = message.split(":")
|
||||
if message[0] == self.username:
|
||||
return COLORS["my_username_color"] + message[0] + ": " + message[1] + COLORS["text_color"]
|
||||
return message[0] + ": " + message[1] + COLORS["text_color"]
|
||||
|
||||
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:
|
||||
os.system("cls")
|
||||
|
||||
def __print_ip(
|
||||
self,
|
||||
ip: str
|
||||
) -> str:
|
||||
return f"IP: " + COLORS["ip_color"] + ip + COLORS["text_color"]
|
||||
|
||||
def __print_username(
|
||||
self,
|
||||
username: str
|
||||
) -> str:
|
||||
return f"USERNAME: " + COLORS["ip_color"] + username + COLORS["username_color"]
|
||||
|
||||
def __print_chat(self, response: list[str]) -> str:
|
||||
for i, msg in enumerate(response["messages"]):
|
||||
actual_message = self._decrypt(msg)
|
||||
if i == 0:
|
||||
for user in response["users_in_chat"]:
|
||||
print(self.__print_ip(user.split(",")[0]))
|
||||
print(self.__print_username(user.split(",")[1]))
|
||||
print(f"\n{self.__print_message(actual_message)}")
|
||||
else:
|
||||
print(f"{self.__print_message(actual_message)}")
|
||||
|
||||
def update_info(self):
|
||||
""" connecting to websocket,
|
||||
wating for updates,
|
||||
updating every 0.05 seconds
|
||||
"""
|
||||
ws = create_connection(f"{self.ws_url}/update")
|
||||
last_try = None
|
||||
while True:
|
||||
try:
|
||||
time.sleep(0.05)
|
||||
response = eval(ws.recv())
|
||||
if last_try == response:
|
||||
continue
|
||||
last_try = response
|
||||
self.__clear_console()
|
||||
if len(last_try["messages"]) > 0:
|
||||
self.__print_chat(
|
||||
response = last_try
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
ws.send(self.close_response)
|
||||
ws.close()
|
||||
quit()
|
||||
except Exception as exc:
|
||||
ws.send(self.close_response)
|
||||
ws.close()
|
||||
print("Something went wrong! ", exc)
|
||||
quit()
|
||||
|
||||
def _validate_keys(self) -> None:
|
||||
self._request_key(self.key_url, self.username)
|
||||
self._remove_keys()
|
||||
|
||||
def run(self):
|
||||
# Running two threads,
|
||||
# One for sending info
|
||||
# Second one for updating info
|
||||
self._validate_keys()
|
||||
threads = [
|
||||
threading.Thread(target=self.send_info),
|
||||
threading.Thread(target=self.update_info)
|
||||
]
|
||||
for th in threads:
|
||||
th.start()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Client(
|
||||
server=input("server ip:\n"),
|
||||
port=int(input("server port: \n")),
|
||||
username=input("username:\n").replace(" ", "").lower()
|
||||
).run()
|
||||
@@ -0,0 +1,8 @@
|
||||
from colorama import Fore
|
||||
|
||||
COLORS = {
|
||||
"text_color": Fore.WHITE,
|
||||
"my_username_color": Fore.MAGENTA,
|
||||
"ip_color": Fore.MAGENTA,
|
||||
"username_color": Fore.GREEN
|
||||
}
|
||||
@@ -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")
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
import rsa
|
||||
import requests
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from cmd_chat.client.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)
|
||||
Reference in New Issue
Block a user