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
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.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
@@ -137,7 +137,7 @@ async def main():
file_size = os.path.getsize(file_path)
chunk_id = 0
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)
async with websockets.connect(WS_RELAY_SERVER) as ws:
msgtype = "announce"
@@ -152,7 +152,6 @@ async def main():
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)
await send_encrypted_msg(ws, k, msg)
await asyncio.sleep(0.05)
chunk_id += 1
if chunk_id > number_of_chunks:
break
@@ -168,35 +167,33 @@ async def main():
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", ""))
bar = None
f = None
i = 0
number_of_chunks = None
while True:
if number_of_chunks is not None and i > number_of_chunks:
bar.index = 100
bar.suffix = '%(percent).1f%% complete'
bar.update()
f.close()
print("")
break
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:
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.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.suffix = '%(percent).1f%% complete'
bar.update()
f.close()
print("")
break