summaryrefslogtreecommitdiff
path: root/day4/task4/client_class.py
blob: 854e5e2b7ee72c40063058451cd8cbc9940617ef (plain)
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
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

        self.server.close()

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