Fix chunk_id checks
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
#!/usr/bin/env python
|
#!/usr/bin/env python
|
||||||
|
|
||||||
import asyncio, websockets
|
import asyncio, websockets
|
||||||
import sys, os, base64, argparse, json, pickle
|
import sys, os, base64, argparse, json, pickle, math
|
||||||
from cryptography.fernet import Fernet
|
from cryptography.fernet import Fernet
|
||||||
from cryptography.hazmat.primitives import hashes
|
from cryptography.hazmat.primitives import hashes
|
||||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||||
@@ -137,7 +137,7 @@ async def main():
|
|||||||
file_size = os.path.getsize(file_path)
|
file_size = os.path.getsize(file_path)
|
||||||
chunk_id = 0
|
chunk_id = 0
|
||||||
chunk_size = 1024 * 512
|
chunk_size = 1024 * 512
|
||||||
number_of_chunks = round(file_size / chunk_size)
|
number_of_chunks = math.ceil(file_size / chunk_size)
|
||||||
WS_RELAY_SERVER = WS_RELAY_SERVER.replace('http', 'ws', 1)
|
WS_RELAY_SERVER = WS_RELAY_SERVER.replace('http', 'ws', 1)
|
||||||
async with websockets.connect(WS_RELAY_SERVER) as ws:
|
async with websockets.connect(WS_RELAY_SERVER) as ws:
|
||||||
msgtype = "announce"
|
msgtype = "announce"
|
||||||
@@ -152,7 +152,6 @@ async def main():
|
|||||||
async for chunk in read_chunks(file_path, chunk_size):
|
async for chunk in read_chunks(file_path, chunk_size):
|
||||||
msg = (msgtype, peer_group_id, role, filename, chunk_size, chunk_id, number_of_chunks, chunk)
|
msg = (msgtype, peer_group_id, role, filename, chunk_size, chunk_id, number_of_chunks, chunk)
|
||||||
await send_encrypted_msg(ws, k, msg)
|
await send_encrypted_msg(ws, k, msg)
|
||||||
await asyncio.sleep(0.05)
|
|
||||||
chunk_id += 1
|
chunk_id += 1
|
||||||
if chunk_id > number_of_chunks:
|
if chunk_id > number_of_chunks:
|
||||||
break
|
break
|
||||||
@@ -168,10 +167,27 @@ async def main():
|
|||||||
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", ""))
|
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", ""))
|
||||||
bar = None
|
bar = None
|
||||||
f = None
|
f = None
|
||||||
i = 0
|
|
||||||
number_of_chunks = None
|
|
||||||
while True:
|
while True:
|
||||||
if number_of_chunks is not None and i > number_of_chunks:
|
message = await ws.recv()
|
||||||
|
message = json.loads(message)
|
||||||
|
payload = message["payload"]
|
||||||
|
msg = pickle.loads(decrypt_chunk(k, payload))
|
||||||
|
chunk_id = msg["chunk_id"]
|
||||||
|
number_of_chunks = msg["number_of_chunks"]
|
||||||
|
if bar is None:
|
||||||
|
bar = Bar('Receiving', max=number_of_chunks, suffix='%(percent).1f%% complete - %(eta_td)s remaining')
|
||||||
|
if f is None:
|
||||||
|
f = open(msg["filename"], "wb")
|
||||||
|
f.write(msg["chunk"])
|
||||||
|
else:
|
||||||
|
f.write(msg["chunk"])
|
||||||
|
msgtype = "proceed"
|
||||||
|
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) # request the next chunk from sender
|
||||||
|
bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining"
|
||||||
|
bar.next()
|
||||||
|
|
||||||
|
if chunk_id >= number_of_chunks-1:
|
||||||
|
# This is the last chunk, exit
|
||||||
bar.index = 100
|
bar.index = 100
|
||||||
bar.suffix = '%(percent).1f%% complete'
|
bar.suffix = '%(percent).1f%% complete'
|
||||||
bar.update()
|
bar.update()
|
||||||
@@ -179,25 +195,6 @@ async def main():
|
|||||||
print("")
|
print("")
|
||||||
break
|
break
|
||||||
|
|
||||||
message = await ws.recv()
|
|
||||||
message = json.loads(message)
|
|
||||||
payload = message["payload"]
|
|
||||||
msg = pickle.loads(decrypt_chunk(k, payload))
|
|
||||||
if bar is None:
|
|
||||||
filename = msg["filename"]
|
|
||||||
number_of_chunks = msg["number_of_chunks"]
|
|
||||||
bar = Bar('Receiving', max=number_of_chunks, suffix='%(percent).1f%% complete - %(eta_td)s remaining')
|
|
||||||
if f is None:
|
|
||||||
f = open(msg["filename"], "wb")
|
|
||||||
f.write(msg["chunk"])
|
|
||||||
else:
|
|
||||||
f.write(msg["chunk"])
|
|
||||||
i += 1
|
|
||||||
msgtype = "proceed"
|
|
||||||
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) # request the next chunk from sender
|
|
||||||
bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining " + str(i) + "/" + str(number_of_chunks) + " msg sent"
|
|
||||||
bar.next()
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
asyncio.run(main())
|
asyncio.run(main())
|
||||||
Reference in New Issue
Block a user