Fix chunk_id checks

This commit is contained in:
SG
2023-05-12 20:01:07 +02:00
parent 6db4849939
commit 603c3b7238

View File

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