-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
156 lines (126 loc) · 6.55 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
from os.path import join, dirname, abspath
from connections.server import ServerConnection
from files.error_handler import ErrorHandler
from connections.connection_management import ServerConnectionsManager
from files.file_manager import FileManager
import os
from session.session_manager import SessionManager
from session.session_manager import get_mac_address
from session.session_manager import get_public_key
if __name__ == "__main__":
choice = input("Do you want to restore your last saved session? (yes/no): ")
# Jeśli użytkownik wybiera przywrócenie sesji
if choice.lower() in ["yes", "y"]:
# Przywracanie zapisanej sesji
session_data = SessionManager.load_session_server()
if session_data:
print("Restoring saved session...")
# Tworzenie serwera z zapisanymi danymi sesji
server_manager = ServerConnectionsManager()
server = ServerConnection(host=session_data["host"], port=session_data["port"], max_connections=1)
server_manager.add_connection(server)
server.start_server()
# Ustawianie zmiennych sesji
client_public_key = session_data["client_public_key"]
server_public_key = session_data["server_public_key"]
aes_key = session_data["aes_key"]
server.restored_server_public_key = server_public_key
server.mac_address = session_data["server_mac_address"]
print("Sending public key to client...")
server.send_public_key_restored()
print("Waiting for client public key...")
client_public_key_received = server.receive_public_key(server.client_socket)
# Sprawdzenie, czy klucz publiczny klienta się zgadza
if get_public_key(client_public_key_received) != get_public_key(client_public_key):
print("[Server] Client's public key does not match. Aborting session.")
server.stop_server()
exit()
file_path = "input2.txt" # Możliwość zmiany pliku
# Szyfrowanie pliku
enc_file_path = FileManager.encrypt_file(file_path, aes_key)
if enc_file_path:
file_data = FileManager.read_file(enc_file_path, 1024, b'/END/')
else:
raise Exception(f"Error during file encryption: {file_path}")
# Sprawdzenie, czy klient jest połączony
if server.client_socket:
# Wysłanie żądania pliku do klienta
server.send_file_request(file_data[1], file_data[2], 1024)
if server.receive_answer():
# Wysłanie zaszyfrowanego pliku do klienta
server.send_file(file_data[0], 1024)
FileManager.delete_file(enc_file_path)
else:
print("[Main] Client declined the file request.")
server.stop_server()
else:
print("[Main] No client connected.")
server.stop_server()
else:
# Jeśli użytkownik nie wybiera przywrócenia sesji, rozpoczynamy nową
server_manager = ServerConnectionsManager()
server = ServerConnection(host='127.0.0.1', port=1234, max_connections=1)
server_manager.add_connection(server)
server.start_server()
server.mac_address = get_mac_address()
try:
file_path = "input.txt"
# Generowanie klucza AES
aes_key = os.urandom(32)
# Szyfrowanie pliku
enc_file_path = FileManager.encrypt_file(file_path, aes_key)
if enc_file_path:
file_data = FileManager.read_file(enc_file_path, 1024, server.message_flags['END'])
else:
raise Exception(f"Error during file encryption: {file_path}")
except Exception as e:
ErrorHandler.error_handling("File_encryption", e)
server.stop_server()
else:
# Sprawdzenie, czy klient jest połączony
if server.client_socket:
try:
print("Sending public key to client...")
server.send_public_key()
# Odbiór klucza publicznego klienta
client_public_key = server.receive_public_key(server.client_socket)
# Szyfrowanie klucza AES kluczem publicznym klienta
encrypted_aes_key = server.encrypt_with_public_key(client_public_key, aes_key)
print("Sending encrypted AES key to client...")
server.client_socket.sendall(encrypted_aes_key)
# Wysłanie żądania pliku do klienta
server.send_file_request(file_data[1], file_data[2], 1024)
if server.receive_answer():
# Wysłanie zaszyfrowanego pliku do klienta
server.send_file(file_data[0], 1024)
FileManager.delete_file(enc_file_path)
# Zapisanie sesji po zakończeniu transferu
save_session_choice = input("Do you want to save this session? (yes/no): ")
if save_session_choice.lower() in ["yes", "y"]:
# Zapisanie sesji
SessionManager.save_session_server(
host=server.host,
port=server.port,
client_public_key=client_public_key,
server_public_key=server.rsa_key_pair.public_key(),
aes_key=aes_key,
file_name=file_path,
server_mac_address=get_mac_address()
)
else:
print("Session not saved.")
else:
print("[Main] Client declined the file request.")
server.stop_server()
FileManager.delete_file(enc_file_path)
except Exception as e:
ErrorHandler.error_handling("Main", e)
server.stop_server()
else:
print("[Main] No client connected.")
server.stop_server()
parent_dir = dirname(abspath(__file__))
priv_dir = join(parent_dir, '.private_keys')
public_dir = join(parent_dir, '.public_keys')
FileManager.delete_file(join(public_dir, 's_public_key.pem'))
FileManager.delete_file(join(priv_dir, 's_private_key.pem'))