-
Notifications
You must be signed in to change notification settings - Fork 0
/
transfer.py
480 lines (365 loc) · 11.7 KB
/
transfer.py
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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
from enum import Enum
import struct
from struct import pack, unpack
from os.path import getsize, basename
from random import random
import utils
from utils import pretty_print
import config.settings as settings
from custom_logging import logger
from logged_exception import LoggedException
class Constant(Enum):
BEG_TX = 0x0000
END_TX = 0x1111
BEG_SIZE = 0x2222
END_SIZE = 0x3333
BEG_COUNT = 0x2020
END_COUNT = 0x3030
BEG_DATA = 0x4444
END_DATA = 0x5555
BEG_HASH = 0x6666
END_HASH = 0x7777
MSG_TYPE_CHUNK = 0x8888
MSG_TYPE_INFOREQ = 0x9999
MSG_TYPE_FILE = 0xBBBB
MSG_TYPE_STRING = 0xCCCC
BEG_CHUNK = 0xDDDD
END_CHUNK = 0xEEEE
BEG_FNAME = 0xAAAA
END_FNAME = 0xFFFF
INFO_HASH_MISMATCH = 0xDEAD
INFO_RECV_OK = 0x600D
REQ_RESEND = 0x5E9D
# -----------------------------------------------------------------------------
# ----------------------------- Transmission ----------------------------------
# -----------------------------------------------------------------------------
"""Convenience function that I'm not sure I should use"""
def fix_str(s):
if type(s) == str:
s = s.encode()
return s
"""
Get a packed version of a two-byte constant (unsigned short), suitable for
network transmission.
get_constant_bytes :: Constant -> Bytes
"""
def constant_to_bytes(constant):
return struct.pack("!H", constant.value)
"""
Convert a string to a packed unsigned short.
get_constant_bytes :: String -> Bytes
"""
def const_string_to_bytes(s):
return struct.pack("!H", Constant[s].value)
"""
Pack an unsigned short.
ushort_to_bytes :: UShort -> Bytes
"""
def ushort_to_bytes(num):
return struct.pack("!H", num)
"""
Pack an unsigned int.
uint_to_bytes :: UInt -> Bytes
This gives us a range of upto 2e32, which means that we can send 4 TB of data
in one request, with 2e32 chunks.
Good enough.
"""
def uint_to_bytes(num):
return struct.pack("!I", num)
"""
Send a constant.
"""
def send_constant(cst, conn):
conn.sendall(constant_to_bytes(cst))
"""
Send a request for something.
"""
def send_req(s, conn):
send_constant(Constant.BEG_TX, conn)
send_constant(Constant.MSG_TYPE_INFOREQ, conn)
send_constant(s, conn)
send_constant(Constant.END_TX, conn)
"""
Send an information opcode.
"""
def send_info(s, conn):
send_constant(Constant.BEG_TX, conn)
send_constant(Constant.MSG_TYPE_INFOREQ, conn)
send_constant(s, conn)
send_constant(Constant.END_TX, conn)
"""
Prepare a byteseq for transmission.
The "grammar" is
BEG_CHUNK
BEG_SIZE <size> END_SIZE
BEG_DATA <data> END_DATA
BEG_HASH <hash> END_HASH
END_CHUNK
An implicit assumption here is that the byteseq passed as an argument is not
larger than the chunk size set in the `settings` module.
"""
def bytes_to_hashed_chunk(s):
if len(s) > settings.CHUNK_SIZE:
logger.error("Too big a byteseq to fit in one chunk")
return
b = bytearray()
# Marker for the beginning of the chunk
b += constant_to_bytes(Constant.BEG_CHUNK)
# "Size block"
b += constant_to_bytes(Constant.BEG_SIZE)
b += uint_to_bytes(len(s))
b += constant_to_bytes(Constant.END_SIZE)
# "Data block"
b += constant_to_bytes(Constant.BEG_DATA)
b += s
b += constant_to_bytes(Constant.END_DATA)
# "Hash block"
b += constant_to_bytes(Constant.BEG_HASH)
b += utils.get_hash(s)
b += constant_to_bytes(Constant.END_HASH)
# Marker for the end of the chunk
b += constant_to_bytes(Constant.END_CHUNK)
# logger.debug("Constructed chunk {} for byteseq {}"
# .format(pretty_print(b), pretty_print(s)))
return bytes(b)
"""
Takes a chunk and sends it over.
"""
def send_chunk(chunk, conn, num_retries=settings.MAX_RETRIES, ab=(0, 0)):
tries_left = num_retries
while tries_left > 0:
conn.sendall(chunk)
if recv_inforeq(conn) != Constant.INFO_RECV_OK:
tries_left -= 1
logger.error(
"Chunk {} of {} apparently received incorrectly "
"({} out of {} retries left)"
.format(*ab, tries_left, num_retries))
else:
logger.debug("Chunk {} of {} sent successfully".format(*ab))
break
"""
The grammar is:
BEG_TX
MSG_TYPE_STRING
BEG_COUNT <number of chunks> END_COUNT
<chunks>
END_TX
"""
def send_string(s, conn, num_retries=settings.MAX_RETRIES):
lst = utils.chunks_of(s, settings.CHUNK_SIZE)
# Mark the beginning of the transaction
send_constant(Constant.BEG_TX, conn)
send_constant(Constant.MSG_TYPE_STRING, conn)
send_constant(Constant.BEG_COUNT, conn)
conn.send(uint_to_bytes(len(lst)))
send_constant(Constant.END_COUNT, conn)
for ix, ch in enumerate([k.encode() for k in lst]):
send_chunk(bytes_to_hashed_chunk(ch),
conn, num_retries=num_retries, ab=(ix+1, len(lst)))
send_constant(Constant.END_TX, conn)
"""
The grammar is:
BEG_TX
MSG_TYPE_FILE
BEG_SIZE <length of file name> END_SIZE
BEG_FNAME <file name> END_FNAME
BEG_COUNT <number of chunks> END_COUNT
<chunks>
END_TX
"""
def send_file(filepath, conn, num_retries=settings.MAX_RETRIES):
flen = getsize(filepath)
num_chunks = flen // settings.FILE_CHUNK_SIZE
num_done = 0
fname = basename(filepath)
# Mark the beginning of the transaction
send_constant(Constant.BEG_TX, conn)
send_constant(Constant.MSG_TYPE_FILE, conn)
send_constant(Constant.BEG_SIZE, conn)
conn.send(uint_to_bytes(len(fname.encode())))
send_constant(Constant.END_SIZE, conn)
send_constant(Constant.BEG_FNAME, conn)
conn.send(fname.encode())
send_constant(Constant.END_FNAME, conn)
send_constant(Constant.BEG_COUNT, conn)
conn.send(uint_to_bytes(num_chunks))
send_constant(Constant.END_COUNT, conn)
with open(filepath, 'rb') as f:
while True:
b = f.read(settings.FILE_CHUNK_SIZE)
if b:
send_chunk(bytes_to_hashed_chunk(b),
conn, num_retries=num_retries,
ab=(num_done, num_chunks))
num_done += 1
else:
break
send_constant(Constant.END_TX, conn)
# -----------------------------------------------------------------------------
# ------------------------------- Reception -----------------------------------
# -----------------------------------------------------------------------------
"""
Convert a packed unsigned short into a Constant.
bytes_to_constant :: Bytes -> Constant
"""
def bytes_to_constant(b):
return Constant((struct.unpack("!H", b))[0])
"""
Unpack an unsigned short.
bytes_to_ushort :: Bytes -> UShort
"""
def bytes_to_ushort(b):
return (struct.unpack("!H", b))[0]
"""
Unpack an unsigned int.
bytes_to_ushort :: Bytes -> UShort
"""
def bytes_to_uint(b):
return (struct.unpack("!I", b))[0]
"""
Return None if parse fails because this is not Haskell and I don't know the
Pythonic way to emulate Maybe.
hashed_chunk_to_bytes :: Bytes -> Maybe Bytes
"""
def hashed_chunk_to_bytes(ch):
pass
"""
Receive a constant, really.
"""
def recv_inforeq(conn):
try:
match_next(Constant.BEG_TX, conn)
match_next(Constant.MSG_TYPE_INFOREQ, conn)
ir = bytes_to_constant(conn.recv(2))
match_next(Constant.END_TX, conn)
return ir
except LoggedException as le:
le.log()
"""
"""
def match_next(opcode, conn):
received_bytes = conn.recv(2)
expected_bytes = constant_to_bytes(opcode)
if received_bytes != expected_bytes:
logger.error("Unexpected opcode: "
"expected {}, received {}"
.format(
pretty_print(expected_bytes),
pretty_print(received_bytes)))
else:
pass # logger.info("Received expected opcode {}".format(opcode))
"""
"""
def consume_till_next(opcode, conn):
expected_bytes = constant_to_bytes(opcode)
(fst, snd) = (expected_bytes[:1], expected_bytes[1:])
skip = bytearray()
done = False
try:
while not done:
rfst = conn.recv(1)
if rfst != fst:
skip += rfst
continue
else:
rsnd = conn.recv(1)
if rsnd != snd:
skip += rsnd
else:
done = True
if len(skip) > 0:
logger.warning("Skipping {} bytes".format(len(skip)))
except LoggedException as le:
le.log()
# logger.info("Received expected opcode {}"
# .format(opcode))
"""
Receive a packed ushort.
"""
def recv_ushort(conn):
return bytes_to_ushort(conn.recv(2))
"""
Receive a packed uint.
"""
def recv_uint(conn):
return bytes_to_uint(conn.recv(4))
"""
Receive a chunk.
"""
def recv_chunk(conn, num_retries):
tries_left = num_retries
while tries_left > 0:
try:
# Assume that the previous send failed, and skip over as much of
# the stream as is required.
consume_till_next(Constant.BEG_CHUNK, conn)
match_next(Constant.BEG_SIZE, conn)
size = recv_uint(conn)
match_next(Constant.END_SIZE, conn)
match_next(Constant.BEG_DATA, conn)
data = conn.recv(size)
match_next(Constant.END_DATA, conn)
match_next(Constant.BEG_HASH, conn)
actual_hash = conn.recv(settings.HASH_LEN)
match_next(Constant.END_HASH, conn)
match_next(Constant.END_CHUNK, conn)
calc_hash = utils.get_hash(data)
if (calc_hash == actual_hash):
logger.debug("Chunk received, hashes match (both {})"
.format(pretty_print(calc_hash)))
send_info(Constant.INFO_RECV_OK, conn)
return data
else:
tries_left -= 1
send_info(Constant.INFO_HASH_MISMATCH, conn)
raise LoggedException("Hash mismatch!")
except LoggedException as le:
tries_left -= 1
le.log()
"""
Receive a string.
"""
def recv_string(conn, num_retries=settings.MAX_RETRIES):
ba = bytearray()
try:
consume_till_next(Constant.BEG_TX, conn)
match_next(Constant.MSG_TYPE_STRING, conn)
match_next(Constant.BEG_COUNT, conn)
num_chunks = recv_uint(conn)
match_next(Constant.END_COUNT, conn)
for i in range(num_chunks):
ba += recv_chunk(conn, num_retries=num_retries)
match_next(Constant.END_TX, conn)
# logger.debug("Received string {}".format(pretty_print(ba)))
return bytes(ba)
except LoggedException as le:
le.log()
"""
Receive a file.
"""
def recv_file(conn, num_retries=settings.MAX_RETRIES):
try:
consume_till_next(Constant.BEG_TX, conn)
match_next(Constant.MSG_TYPE_FILE, conn)
match_next(Constant.BEG_SIZE, conn)
fname_size = recv_uint(conn)
logger.debug("File name is {} bytes long".format(fname_size))
match_next(Constant.END_SIZE, conn)
match_next(Constant.BEG_FNAME, conn)
fname = conn.recv(fname_size).decode()
logger.debug("Awaiting file " + fname)
match_next(Constant.END_FNAME, conn)
match_next(Constant.BEG_COUNT, conn)
num_chunks = recv_uint(conn)
num_left = num_chunks
match_next(Constant.END_COUNT, conn)
with open('./tmp/recv/{}'.format(fname), 'wb') as f:
while num_left >= 0:
f.write(recv_chunk(conn, num_retries=num_retries))
num_left -= 1
match_next(Constant.END_TX, conn)
# logger.debug("Received string {}".format(pretty_print(ba)))
# return bytes(ba)
except LoggedException as le:
le.log()