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
|
import socket
from time import sleep
import struct
import json
import shelve
try:
from tqdm import tqdm
TQDM = True
except ImportError:
print('Для вывода прогресса загрузки файлов установите библиотеку tqdm')
TQDM = False
DATATYPE_SIZE = 4 # int
CHUNK_SIZE = 4096000
HOST, PORT = ADDR = 'localhost', 6000
class Client:
def __init__(self, server_host, server_port):
self.server = socket.socket()
self.db = shelve.open('backdoor.db')
address = server_host, server_port
cnt = 0
while cnt < 10:
try:
cnt += 1
self.server.connect(address)
except ConnectionRefusedError:
print('Не удалось подключиться к серверу. Следующая попытка через 3 секунды...')
sleep(3)
else:
break
else:
print('Не удалось подключиться к серверу через 10 попыток.\nОстанавливаемся...')
self.server.close()
def start(self):
while True:
try:
data = input('Введите команду:\n')
self.handle_input(data)
except ValueError as err:
print(err)
except KeyboardInterrupt:
break
def handle_input(self, data):
cmd, args = data.split(' ', maxsplit=1)
if cmd == 'get':
self.download_file(args)
elif cmd == 'cp':
paths = args.split()
if len(paths) != 2:
raise ValueError('Неверное количество путей в параметрах.')
src, dest = paths
# Если файл не присутствует в БД, то предварительно его загружаем
if src not in self.db:
self.download_file(src)
# При любых ошибках при загрузке файла download_file выкинет исключение и исполнение
# не дойдет до этого блока, поэтому можно не проверять наличие src в БД
with open(dest, 'wb') as f:
f.write(self.db[src])
else:
raise ValueError('Команда не распознана')
def download_file(self, path):
response = self.make_get_request(path)
if response['status'] == 0:
print(f'Загружаем файл размером {response["size"]} байт...')
file_data = self.receive_data(int(response['size']), use_tqdm=True)
self.db[path] = file_data
print('Загрузка завершена')
elif response['status'] == 1:
raise ValueError(response['reason'])
elif response['status'] == 2:
raise ValueError(f'Неизвестная ошибка: {response["reason"]}')
def make_get_request(self, path):
packet = Client.pack_str(json.dumps({'method': 'get', 'path': path}))
self.server.sendall(packet)
return Client.parse_response(self.receive_data(DATATYPE_SIZE))
def receive_data(self, data_size, use_tqdm=False):
data = b''
# С копипастой надо бы что-то сделать, но я не знаю что
if not use_tqdm or not TQDM:
while len(data) < data_size:
(packet_size,) = struct.unpack('I', self.server.recv(DATATYPE_SIZE))
data += self.server.recv(packet_size)
if not data:
raise ConnectionError('Connection was closed')
elif TQDM:
for count in tqdm(range(0, data_size, CHUNK_SIZE)):
(packet_size,) = struct.unpack('I', self.server.recv(DATATYPE_SIZE))
data += self.server.recv(packet_size)
if not data:
raise ConnectionError('Connection was closed')
return data
@staticmethod
def pack_str(data):
encoded = data.encode()
return struct.pack('I', len(encoded)) + encoded
@staticmethod
def parse_response(data):
return json.loads(data.decode())
client = Client(HOST, PORT)
client.start()
|
