137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
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
|
|
|
|
init()
|
|
|
|
OS = "Windows"
|
|
if "Linux" in str(platform.platform()):
|
|
OS = "Linux"
|
|
|
|
class Client:
|
|
|
|
def _key_gen(self) -> None:
|
|
(pubkey, privkey) = rsa.newkeys(512)
|
|
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.port = input("server port: \n")
|
|
self.username = input("your username: \n")
|
|
# 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
|
|
self.privkey = None
|
|
self.symetric_key = None
|
|
self.fernet = None
|
|
|
|
def send_info(self):
|
|
while True:
|
|
user_input = input("You're message: ")
|
|
message = f'{self.username}: {user_input}'
|
|
requests.post(self.talk_url, data={
|
|
"text": self.fernet.encrypt(message.encode()),
|
|
"username": self.username
|
|
})
|
|
|
|
def print_message(self, message: str) -> str:
|
|
message = message.split(":")
|
|
if message[0] == self.username:
|
|
return Fore.MAGENTA + message[0] + ": " + message[1] + Fore.WHITE
|
|
else:
|
|
return message[0] + ": " + message[1] + Fore.WHITE
|
|
|
|
def update_info(self):
|
|
ws = create_connection(f"{self.ws_url}/update")
|
|
last_try = None
|
|
while True:
|
|
time.sleep(0.05)
|
|
#r = requests.post(self.info_url)
|
|
ws.send(payload="")
|
|
r = ws.recv()
|
|
if last_try == r.json():
|
|
continue
|
|
last_try = r.json()
|
|
# For windows clear command its cls
|
|
# For linux clear command its clear
|
|
if OS == "Linux":
|
|
os.system("clear")
|
|
else:
|
|
os.system("cls")
|
|
if len(last_try['status']) > 0:
|
|
for i, msg in enumerate(last_try["status"]):
|
|
actual_message = self.fernet.decrypt(
|
|
msg.encode()
|
|
).decode("utf-8")
|
|
if i == 0:
|
|
users = last_try["users_in_chat"]
|
|
for user in users:
|
|
ip = user.split(",")[0]
|
|
username = user.split(",")[1]
|
|
print("IP:", Fore.MAGENTA + ip + Fore.WHITE)
|
|
print("USERNAME: ", Fore.GREEN + username + Fore.WHITE)
|
|
print()
|
|
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._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):
|
|
# 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()()
|