1 files changed, 47 insertions, 29 deletions

diff --git a/day4/task4/client.py b/day4/task4/client.py
index 33ea084..bc4a40f 100644
--- a/day4/task4/client.py
+++ b/day4/task4/client.py
@@ -4,40 +4,60 @@ import struct
 import json
 import shelve
 
+try:
+    from tqdm import tqdm
+    TQDM = True
+except ImportError:
+    print('Для вывода прогресса загрузки файлов установите библиотеку tqdm')
+    TQDM = False
+
 
 def pack_data(data):
     encoded = data.encode()
     return struct.pack('I', len(encoded)) + encoded
 
 
-def prepare_data(data):
+def handle_input(data):
     cmd, args = data.split(' ', maxsplit=1)
 
     if cmd == 'get':
-        return json.dumps({'method': 'get', 'path': args})
+        make_get_request(args)
 
     elif cmd == 'cp':
         paths = args.split()
         if len(paths) != 2:
             raise ValueError('Неверное количество путей в параметрах.')
-        src, dest = args.split()
+
+        src, dest = paths
+
+        # Если файл не присутствует в БД, то предварительно его загружаем
+        status = 0
         if src not in db:
-            return json.dumps({'method': 'get_cp', 'path': src})
-        with open(dest, 'wb') as f:
-            f.write(db[src])
+            status = make_get_request(src)
+
+        if status == 0:
+            with open(dest, 'wb') as f:
+                f.write(db[src])
 
     else:
         raise ValueError('Команда не распознана')
 
 
+def make_get_request(path):
+    packet = pack_data(json.dumps({'method': 'get', 'path': path}))
+    sock.sendall(packet)
+    response = get_response()
+    return handle_response(response)
+
+
 def parse_response(data):
     return json.loads(data.decode())
 
 
-def get_response(connection):
+def get_response():
     buffer = b''
     while len(buffer) < DATATYPE_SIZE:
-        data = connection.recv(DATATYPE_SIZE - len(buffer))
+        data = sock.recv(DATATYPE_SIZE - len(buffer))
         if not data:
             raise ConnectionError('Connection was closed')
 
@@ -48,25 +68,31 @@ def get_response(connection):
         raise ConnectionError('Connection was closed')
 
     (size,) = struct.unpack('I', size_packed)
-    data = connection.recv(size - len(buffer))
+    data = sock.recv(size - len(buffer))
     return parse_response(buffer + data)
 
 
-def download_file(filesize, connection):
+def download_file(filesize):
     data = b''
-    count = 0
-    while count < filesize:
-        (size,) = struct.unpack('I', connection.recv(DATATYPE_SIZE))
-        data += connection.recv(size)
-        count += size
+
+    if TQDM:
+        for count in tqdm(range(0, filesize, CHUNK_SIZE)):
+            (size,) = struct.unpack('I', sock.recv(DATATYPE_SIZE))
+            data += sock.recv(size)
+    else:
+        count = 0
+        while count < filesize:
+            (size,) = struct.unpack('I', sock.recv(DATATYPE_SIZE))
+            data += sock.recv(size)
+            count += size
 
     return data
 
 
-def handle_response(data, connection):
+def handle_response(data):
     if data['status'] == 0:
         print(f'Загружаем файл размером {data["size"]} байт...')
-        file_data = download_file(data['size'], connection)
+        file_data = download_file(data['size'])
         db[data['path']] = file_data
         print('Загрузка завершена')
 
@@ -77,6 +103,8 @@ def handle_response(data, connection):
         print('Неизвестная ошибка:')
         print(data['reason'])
 
+    return data['status']
+
 
 def try_connect():
     cnt = 0
@@ -101,6 +129,7 @@ def try_connect():
 
 HOST, PORT = ADDR = 'localhost', 6000
 DATATYPE_SIZE = 4  # int
+CHUNK_SIZE = 4096000
 
 sock = try_connect()
 if sock is None:
@@ -112,20 +141,9 @@ db = shelve.open('backdoor.db')
 while True:
     try:
         data = input('Введите команду:\n')
-        data = prepare_data(data)
-        if data is None:
-            continue
+        handle_input(data)
 
     except KeyboardInterrupt:
         break
 
-    except ValueError as err:
-        print(err)
-
-    else:
-        packet = pack_data(data)
-        sock.sendall(packet)
-        response = get_response(sock)
-        handle_response(response, sock)
-
 sock.close()