Skip to content
This repository was archived by the owner on Jul 30, 2024. It is now read-only.

Commit 87c9204

Browse files
aniko33aniko33
authored
Update main.py
1 parent f7f1108 commit 87c9204

File tree

1 file changed

+211
-102
lines changed

1 file changed

+211
-102
lines changed

‎client/main.py‎

Lines changed: 211 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -8,81 +8,213 @@
88

99
from rich import print
1010
from rich.console import Console
11-
from rich.tree import Tree
11+
from rich.tree import Tree
1212
from rich.align import Align
1313

14+
global buffer
15+
16+
# Vars
1417
console = Console()
1518
s = socket.socket()
16-
17-
#var
1819
buffer = 1024
1920

21+
22+
"""
23+
def cls:
24+
Clear terminal screen
25+
"""
26+
27+
2028
def cls():
21-
os.system('cls' if os.name == 'nt' else 'clear')
29+
if os.name == 'nt':
30+
os.system('cls')
2231

23-
class Chat:
24-
def __init__(self, priv_key, pub_key) -> None:
25-
self.priv_key = priv_key
26-
self.pub_key = pub_key
27-
def send(self, msg:str):
28-
s.send(rsa.encrypt(msg.encode(), self.pub_key))
29-
def recv(self, buffer: int):
30-
msg = s.recv(buffer)
31-
return rsa.decrypt(msg, self.priv_key).decode()
32+
else:
33+
os.system("clear")
3234

33-
class Main:
34-
def __init__(self) -> None:
35-
pass
3635

37-
class UI:
38-
def __init__(self) -> None:
39-
pass
40-
def random_color(self):
41-
#color list: https://rich.readthedocs.io/en/stable/appendix/colors.html
42-
colors = ["red", "blue", "cyan", "yellow", "magenta", "green", "purple", "violet", "gold"]
43-
44-
#random choice from array
45-
color = random.choice(colors)
46-
return color
47-
def banner(self):
48-
print(Align("""[purple]
49-
__________________
50-
/ ___/ __/ __/ ___/
51-
/ /___\ \/ _// /__
52-
\___/___/___/\___/
53-
""","center"))
54-
def get_server(self):
55-
server_ip = console.input("[b]Insert server address[/b] [purple]>>[/purple] ")
56-
server_port = int(console.input("[b]Insert server port[/b] [purple]>>[/purple] "))
57-
return server_ip, server_port
58-
59-
def get_username(self):
60-
username = console.input("[b]Insert your username[/b] [purple]>>[/purple] ")
61-
#random color set
62-
color = self.random_color()
63-
username_styled = f"<[{color}]{username}[/{color}]>"
64-
return username, username_styled
65-
66-
def start(self):
67-
self.banner()
68-
ip, port = self.get_server()
69-
username, username_styled = self.get_username()
70-
return ip, port, username, username_styled
36+
class API:
37+
"""
38+
class Chat:
39+
arg: private_key, public_key
40+
41+
def send (message: str):
42+
Allows you to send an encrypted message to the server
43+
44+
def recv (buffer: int):
45+
Receives a message from the encrypted server
46+
47+
class Load_keys:
48+
arg: private_key, public_key
49+
50+
def private:
51+
Load the private key
52+
53+
def public:
54+
Load the public key
55+
56+
def load_all:
57+
Load the private and public key
58+
"""
59+
class Chat:
60+
def __init__(self, priv_key, pub_key) -> None:
61+
self.priv_key = priv_key
62+
self.pub_key = pub_key
63+
64+
def send(self, msg: str):
65+
s.send(rsa.encrypt(msg.encode(), self.pub_key))
66+
67+
def recv(self, buffer: int):
68+
msg = s.recv(buffer)
69+
return rsa.decrypt(msg, self.priv_key).decode()
7170

7271
class Load_keys:
7372
def __init__(self, pub_key, priv_key) -> None:
7473
self.pub_key = pub_key
7574
self.priv_key = priv_key
76-
75+
7776
def private(self):
7877
return rsa.PrivateKey.load_pkcs1(self.priv_key)
7978

8079
def public(self):
8180
return rsa.PublicKey.load_pkcs1(self.pub_key)
8281

8382
def load_all(self):
84-
return rsa.PrivateKey.load_pkcs1(self.priv_key), rsa.PublicKey.load_pkcs1(self.pub_key)
83+
return rsa.PrivateKey.load_pkcs1(
84+
self.priv_key), rsa.PublicKey.load_pkcs1(self.pub_key)
85+
86+
87+
class UI:
88+
"""
89+
def random_color:
90+
It will take a random color and be entered as the username color
91+
92+
def banner:
93+
The ASCII banner that will be centered
94+
95+
def get_server:
96+
Get the server ip and port
97+
98+
def get_username:
99+
Get the username from the input will send it to the server
100+
101+
def run:
102+
Where the code will start
103+
"""
104+
105+
def __init__(self) -> None:
106+
pass
107+
108+
def random_color(self):
109+
# Color list:
110+
# https://rich.readthedocs.io/en/stable/appendix/colors.html
111+
colors = [
112+
"red",
113+
"blue",
114+
"cyan",
115+
"yellow",
116+
"magenta",
117+
"green",
118+
"purple",
119+
"violet",
120+
"gold"]
121+
122+
# Random choice from array
123+
color = random.choice(colors)
124+
return color
125+
126+
def banner(self):
127+
print(Align(r"""[purple]
128+
_____ _____ ______ _____
129+
/ ____|/ ____| ____/ ____|
130+
| | | (___ | |__ | |
131+
| | \___ \| __|| |
132+
| |____ ____) | |___| |____
133+
\_____|_____/|______\_____|
134+
""", "center"))
85135

136+
def get_server(self):
137+
server_ip = console.input(
138+
"[b]Insert server address[/b] [purple]>>[/purple] ")
139+
server_port = int(
140+
console.input("[b]Insert server port[/b] [purple]>>[/purple] "))
141+
return server_ip, server_port
142+
143+
def get_username(self):
144+
username = console.input(
145+
"[b]Insert your username[/b] [purple]>>[/purple] ")
146+
# Random color set
147+
color = self.random_color()
148+
username_styled = f"<[{color}]{username}[/{color}]>"
149+
return username, username_styled
150+
151+
def start(self):
152+
self.banner()
153+
ip, port = self.get_server()
154+
username, username_styled = self.get_username()
155+
return ip, port, username, username_styled
156+
157+
158+
class Chat:
159+
"""
160+
class chat_api: RSA custom socket
161+
162+
var (str) username_syled: (<Username>) {
163+
It is the username but compliant for sending via socket
164+
making all the username available
165+
}
166+
167+
fun receive: Receive messages via socket
168+
fun write: Where the user will put the message, it will be encrypted
169+
fun run: Where the code will start
170+
"""
171+
172+
def __init__(self, chat_api, username_styled) -> None:
173+
self.chat_api = chat_api
174+
self.username_styled = username_styled
175+
176+
def receive(self):
177+
while True:
178+
print(self.chat_api.recv(buffer))
179+
180+
def write(self):
181+
while True:
182+
msg = input()
183+
184+
# Remove line up
185+
sys.stdout.write("\033[F")
186+
187+
if (len(msg.strip()) > 0):
188+
self.chat_api.send(self.username_styled + " " + msg)
189+
print("<[green][i]You[/i][/green]> " + msg)
190+
191+
def run(self):
192+
receive_process = threading.Thread(target=self.receive)
193+
receive_process.start()
194+
write_process = threading.Thread(target=self.write)
195+
write_process.start()
196+
197+
198+
class Main:
199+
"""
200+
def connect:
201+
It allows you to connect to the server, you will receive the RSA keys
202+
203+
def send_username:
204+
Send yours username to server
205+
206+
def get_buffer:
207+
Get from server buffer - RSA size and message size
208+
209+
def get_welcome_message:
210+
Get welcome message from server
211+
212+
def run:
213+
Where the code will start
214+
"""
215+
216+
def __init__(self) -> None:
217+
pass
86218

87219
def connect(self):
88220
while True:
@@ -93,70 +225,47 @@ def connect(self):
93225
print("[red]ERROR[/red]: Connection refused")
94226
UI = Main.UI()
95227
self.ip, self.port = UI.get_server()
96-
self.username = UI.get_username()
228+
self.username, self.username_styled = UI.get_username()
97229
else:
98230
break
99-
231+
232+
self.send_username(self.username)
233+
100234
buffer = self.get_buffer()
101235

102-
#decompress keys with zlib
103-
private_key = zlib.decompress(s.recv(buffer))
236+
# Decompress keys with zlib
104237
public_key = zlib.decompress(s.recv(buffer))
105-
106-
#key load
107-
client_key = self.Load_keys(public_key, private_key)
238+
private_key = zlib.decompress(s.recv(buffer))
239+
240+
# Key load
241+
client_key = API.Load_keys(public_key, private_key)
108242
self.private_key, self.public_key = client_key.load_all()
109-
110-
#API
111-
self.chat_sock = Chat(self.private_key, self.public_key)
112243

113-
#check username existing
114-
self.check_username(self.username)
115-
244+
# API
245+
self.chat_api = API.Chat(self.private_key, self.public_key)
246+
247+
# Welcome screen
116248
cls()
117-
print(Align(self.chat_sock.recv(buffer),"center"))
118-
119-
def check_username(self, username: str):
249+
self.get_welcome_message()
250+
251+
chat = Chat(self.chat_api, self.username_styled)
252+
chat.run()
253+
254+
def send_username(self, username: str):
120255
s.send(username.encode())
121-
status_code = s.recv(buffer)
122-
print(status_code)
123-
status_code = status_code.decode()
124-
if status_code == "False":
125-
s.close()
126-
print("[red]ERROR[/red]: This user exist")
127-
quit()
128256

129257
def get_buffer(self):
130258
return int(s.recv(buffer).decode())
131259

132-
def chat(self):
133-
def receive():
134-
while True:
135-
print(self.chat_sock.recv(buffer))
136-
def write():
137-
while True:
138-
msg = input()
139-
140-
#remove input
141-
sys.stdout.write("\033[F")
142-
143-
if (len(msg.strip()) > 0):
144-
#switch - commands
145-
match msg:
146-
case _:
147-
self.chat_sock.send(self.username_styled+" "+msg)
148-
print("<[green][i]You[i][/green]> "+msg)
149-
receive_process = threading.Thread(target=receive)
150-
receive_process.start()
151-
write_process = threading.Thread(target=write)
152-
write_process.start()
260+
def get_welcome_message(self):
261+
print(Align(self.chat_api.recv(1024), "center"))
153262

154263
def run(self):
155-
UI = Main.UI()
156-
self.ip, self.port, self.username, self.username_styled = UI.start()
264+
ui = UI()
265+
self.ip, self.port, self.username, self.username_styled = ui.start()
157266
self.connect()
158-
self.chat()
159-
267+
268+
160269
if __name__ == "__main__":
161-
main_start = Main()
162-
main_start.run()
270+
main = Main()
271+
main.run()

0 commit comments

Comments
 (0)