Code refactoring
This commit is contained in:
mirai
@@ -0,0 +1,48 @@
|
||||
import asyncio
|
||||
|
||||
from cmd_chat.server.server import app
|
||||
from cmd_chat.client.client import Client
|
||||
|
||||
|
||||
async def run_server(
|
||||
ip: str,
|
||||
port: int
|
||||
) -> None:
|
||||
app.run(
|
||||
host=ip,
|
||||
port=port,
|
||||
dev=False
|
||||
)
|
||||
|
||||
|
||||
async def run_client(
|
||||
username: str,
|
||||
server: str,
|
||||
port: int
|
||||
) -> None:
|
||||
Client(
|
||||
server = server,
|
||||
port = port,
|
||||
username = username
|
||||
).run()
|
||||
|
||||
|
||||
async def run() -> None:
|
||||
action: int = int(input("Choose action:\n1. Run server\n2. Run client\nAction: "))
|
||||
if action == 1:
|
||||
await run_server(
|
||||
input("IP: "),
|
||||
int(input("PORT: "))
|
||||
)
|
||||
if action == 2:
|
||||
await run_client(
|
||||
input("USERNAME: "),
|
||||
input("IP: "),
|
||||
int(input("PORT: "))
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(
|
||||
run()
|
||||
)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,5 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
class Message(BaseModel):
|
||||
|
||||
message: str
|
||||
@@ -0,0 +1,68 @@
|
||||
import asyncio
|
||||
|
||||
import rsa
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from sanic.response import HTTPResponse
|
||||
from sanic import Sanic, Request, response, Websocket
|
||||
|
||||
from cmd_chat.server.models import Message
|
||||
from cmd_chat.server.services import (
|
||||
_get_bytes_and_serialize,
|
||||
_check_ws_for_close_status,
|
||||
_generate_new_message,
|
||||
_generate_update_payload
|
||||
)
|
||||
|
||||
|
||||
app = Sanic("app")
|
||||
app.config.OAS = False
|
||||
|
||||
|
||||
# Message structure is:
|
||||
# [username: message, ...]
|
||||
MESSAGES_MEMORY_DB: list[Message] = []
|
||||
|
||||
|
||||
# Users structure is
|
||||
# {Ip, Username: Public key}
|
||||
USERS: dict[str, str] = {}
|
||||
PUBLIC_KEY = Fernet.generate_key()
|
||||
|
||||
|
||||
@app.websocket("/talk")
|
||||
async def talk_ws_view(request: Request, ws: Websocket) -> HTTPResponse:
|
||||
while True:
|
||||
serialized_message: dict = await _get_bytes_and_serialize(ws)
|
||||
await _check_ws_for_close_status(
|
||||
serialized_message,
|
||||
ws
|
||||
)
|
||||
new_message = await _generate_new_message(
|
||||
serialized_message.get("text")
|
||||
)
|
||||
MESSAGES_MEMORY_DB.append(new_message)
|
||||
await ws.send(
|
||||
str({"status": "ok"})
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
|
||||
@app.websocket("/update")
|
||||
async def update_ws_view(request: Request, ws: Websocket) -> HTTPResponse:
|
||||
while True:
|
||||
payload = await _generate_update_payload(
|
||||
MESSAGES_MEMORY_DB,
|
||||
USERS
|
||||
)
|
||||
await ws.send(payload.encode())
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
|
||||
@app.route('/get_key', methods=['GET', 'POST'])
|
||||
async def get_key_view(request: Request) -> HTTPResponse:
|
||||
public_key = rsa.PublicKey.load_pkcs1(request.form.get('pubkey'))
|
||||
encrypted_data = rsa.encrypt(PUBLIC_KEY, public_key)
|
||||
if request.ip not in USERS:
|
||||
USERS[f"{request.ip}, {request.form.get('username')}"] = PUBLIC_KEY
|
||||
return response.raw(encrypted_data)
|
||||
@@ -0,0 +1,35 @@
|
||||
from sanic import Websocket
|
||||
from cmd_chat.server.models import Message
|
||||
|
||||
|
||||
async def _get_bytes_and_serialize(
|
||||
ws: Websocket
|
||||
) -> dict:
|
||||
return eval(await ws.recv())
|
||||
|
||||
|
||||
async def _check_ws_for_close_status(
|
||||
response: dict,
|
||||
ws: Websocket
|
||||
) -> None:
|
||||
if "action" in response.keys():
|
||||
if response["action"] == "close":
|
||||
await ws.close()
|
||||
|
||||
|
||||
async def _generate_new_message(
|
||||
message: str
|
||||
) -> Message:
|
||||
return Message(message = message)
|
||||
|
||||
|
||||
async def _generate_update_payload(
|
||||
memory_msgs: list[str],
|
||||
users_structure: dict
|
||||
) -> str:
|
||||
return str({
|
||||
"messages": [i.message for i in memory_msgs],
|
||||
"users_in_chat": list(users_structure.keys())
|
||||
})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user