diff --git a/transmat/transmat.py b/transmat/transmat.py index e363008..c363fb2 100755 --- a/transmat/transmat.py +++ b/transmat/transmat.py @@ -9,16 +9,24 @@ from progress.bar import Bar import diceware -VERSION_NUMBER="0.1.8" +VERSION_NUMBER="0.2.0" +PROTOCOL_VERSION="0.2" SCRIPT_BASENAME = os.path.basename(sys.argv[0]) -async def read_chunks(filename, chunk_size = 1024): - with open(filename, "rb") as f: - while True: - chunk = f.read(chunk_size) - if not chunk: - break - yield chunk +#async def read_chunks(filename, chunk_size = 1024): +# with open(filename, "rb") as f: +# while True: +# chunk = f.read(chunk_size) +# if not chunk: +# break +# yield chunk + + +def chunk_from_file(file, chunk_id, chunk_size): + position = chunk_id * chunk_size + file.seek(position) + chunk = file.read(chunk_size) + return chunk def derive_key_from_password(password): @@ -68,9 +76,7 @@ async def send_msg(ws, msg): async def send_encrypted_msg(ws, k, data): ( - msgtype, peer_group_id, - role, filename, chunk_size, chunk_id, @@ -86,9 +92,7 @@ async def send_encrypted_msg(ws, k, data): "chunk": chunk } msg = { - "msgtype": msgtype, "peer_group_id": peer_group_id, - "role": role, "payload": encrypt_chunk(k, pickle.dumps(payload)) } await send_msg(ws, json.dumps(msg)) @@ -125,7 +129,7 @@ async def main(): WS_RELAY_SERVER = args.server server_part = f'--server {WS_RELAY_SERVER}' if args.version: - print(f"{SCRIPT_BASENAME} ver {VERSION_NUMBER}") + print(f"{SCRIPT_BASENAME} ver. {VERSION_NUMBER}, protocol ver. {PROTOCOL_VERSION}") sys.exit(0) k = derive_key_from_password(password) @@ -135,60 +139,75 @@ async def main(): print('Run the following command on the remote party:\n\n', send_part, passwd_part, server_part, "\n") filename = os.path.basename(file_path) file_size = os.path.getsize(file_path) - chunk_id = 0 + file = open(file_path, 'rb') chunk_size = 1024 * 512 + chunk_id = "" + chunk = "" 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" - await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, filename, "", "", number_of_chunks, "")) + # Announce ourself to the relay + await send_encrypted_msg(ws, k, (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk)) + bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') while True: + # Wait for messages from receiver message = await ws.recv() message = json.loads(message) - if message["msgtype"] == "announce" and message["peer_group_id"] == peer_group_id: - break - bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') - msgtype = "data" - 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) + # If message is not meant for us - skip it + if 'peer_group_id' not in message or message["peer_group_id"] != peer_group_id: + pass + payload = pickle.loads(decrypt_chunk(k, message["payload"])) + chunk_id = payload["chunk_id"] + chunk = chunk_from_file(file, chunk_id, chunk_size) + #bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') + msg = (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk) await send_encrypted_msg(ws, k, msg) + #bar.index = chunk_id + bar.next() chunk_id += 1 if chunk_id > number_of_chunks: - break - proceed = await ws.recv() - bar.next() + break + bar.suffix = '%(percent).1f%% complete' bar.update() print("") if role =='receive': async with websockets.connect(WS_RELAY_SERVER) as ws: - msgtype = "announce" - await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) + filename = "" + chunk_size = "" + chunk_id = 0 + number_of_chunks = "" + chunk = "" bar = None f = None while True: + # Announce ourself to the relay + await send_encrypted_msg(ws, k, (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk)) + # Wait for message from sender message = await ws.recv() message = json.loads(message) + # If message is not meant for us - skip it + if 'peer_group_id' not in message or message["peer_group_id"] != peer_group_id: + pass payload = message["payload"] msg = pickle.loads(decrypt_chunk(k, payload)) chunk_id = msg["chunk_id"] number_of_chunks = msg["number_of_chunks"] filename = msg["filename"] if bar is None: - bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') + bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') if f is None: f = open(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 + chunk_id += 1 bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining" bar.next() - - if chunk_id >= number_of_chunks-1: + if chunk_id >= number_of_chunks: # This was the last chunk, exit + await send_encrypted_msg(ws, k, (peer_group_id, "", "", chunk_id, "", "")) f.close() bar.suffix = '%(percent).1f%% complete' bar.update()