Немного изменил структуру клиента и переписал его с помощью ООП

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()