check_tty=False

transmat/transmat.py

@@ -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
@@ -95,7 +95,7 @@ async def send_encrypted_msg(ws, k, data):
 
 
 async def main():
-    WS_RELAY_SERVER = "wss://transmat.exocortex.ru"
+    WS_RELAY_SERVER = "wss://transmat.exocortex.ru/ws"
     parser = argparse.ArgumentParser()
     arg_group = parser.add_mutually_exclusive_group(required=True)
     arg_group.add_argument('--receive', '--recv', action='store_true', help='Receive a file from the remote party (mutually exclusive with --send and --relay)')
@@ -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"
@@ -147,12 +147,11 @@ async def main():
                 message = json.loads(message)
                 if message["msgtype"] == "announce" and message["peer_group_id"] == peer_group_id:
                     break
-            bar = Bar('Transferring', max=number_of_chunks,  suffix='%(percent).1f%% complete - %(eta_td)s remaining')
+            bar = Bar('Transferring', 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)
                 await send_encrypted_msg(ws, k, msg)
-                await asyncio.sleep(0.05)
                 chunk_id += 1
                 if chunk_id > number_of_chunks:
                     break
@@ -168,16 +167,15 @@ async def main():
             await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", ""))
             bar = None
             f = None
-            i = 1
             while True:
                 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')
+                    bar = Bar('Receiving', max=number_of_chunks,  check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining')
                 if f is None:
                     f = open(msg["filename"], "wb")
                     f.write(msg["chunk"])
@@ -185,15 +183,17 @@ async def main():
                     f.write(msg["chunk"])
                 msgtype = "proceed"
                 await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) # request the next chunk from sender
-                i += 1
+                bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining"
                 bar.next()
-                if i > msg["number_of_chunks"] + 1:
-                    bar.index = 100
+
+                if chunk_id >= number_of_chunks-1:
+                    # This is the last chunk, exit
+                    f.close()
                     bar.suffix = '%(percent).1f%% complete'
                     bar.update()
-                    f.close()
-                    print("")
+                    bar.finish()
                     break
+ 
                 
 
 asyncio.run(main())