diff options
Diffstat (limited to 'day4')
| -rw-r--r-- | day4/task4/client_class.py | 132 |
1 files changed, 132 insertions, 0 deletions
diff --git a/day4/task4/client_class.py b/day4/task4/client_class.py new file mode 100644 index 0000000..0925094 --- /dev/null +++ b/day4/task4/client_class.py @@ -0,0 +1,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() |