diff --git a/src/wormhole/scripts/cmd_receive_file.py b/src/wormhole/scripts/cmd_receive_file.py index b51371c..aa73980 100644 --- a/src/wormhole/scripts/cmd_receive_file.py +++ b/src/wormhole/scripts/cmd_receive_file.py @@ -3,6 +3,7 @@ import sys, os, json from nacl.secret import SecretBox from wormhole.blocking.transcribe import Receiver, WrongPasswordError from wormhole.blocking.transit import TransitReceiver +from .progress import start_progress, update_progress, finish_progress APPID = "lothar.com/wormhole/file-xfer" @@ -39,17 +40,25 @@ def receive_file(so): transit_receiver.add_their_direct_hints(tdata["direct_connection_hints"]) transit_receiver.add_their_relay_hints(tdata["relay_connection_hints"]) skt = transit_receiver.establish_connection() - print("Receiving %d bytes.." % filesize) + + print("Receiving %d bytes for '%s'.." % (filesize, filename)) + encrypted = b"" + next_update = start_progress(encrypted_filesize) while len(encrypted) < encrypted_filesize: more = skt.recv(encrypted_filesize - len(encrypted)) if not more: + print() print("Connection dropped before full file received") print("got %d bytes, wanted %d" % (len(encrypted), encrypted_filesize)) return 1 encrypted += more + next_update = update_progress(next_update, len(encrypted), + encrypted_filesize) + finish_progress(encrypted_filesize) assert len(encrypted) == encrypted_filesize + print("Decrypting..") decrypted = SecretBox(xfer_key).decrypt(encrypted) # only write to the current directory, and never overwrite anything diff --git a/src/wormhole/scripts/cmd_send_file.py b/src/wormhole/scripts/cmd_send_file.py index 1188650..62201cf 100644 --- a/src/wormhole/scripts/cmd_send_file.py +++ b/src/wormhole/scripts/cmd_send_file.py @@ -3,6 +3,7 @@ import os, sys, json from nacl.secret import SecretBox from wormhole.blocking.transcribe import Initiator, WrongPasswordError from wormhole.blocking.transit import TransitSender +from .progress import start_progress, update_progress, finish_progress APPID = "lothar.com/wormhole/file-xfer" @@ -38,6 +39,8 @@ def send_file(so): #print("them: %r" % (them_d,)) xfer_key = i.derive_key(APPID+"/xfer-key", SecretBox.KEY_SIZE) + print("Encrypting %d bytes.." % filesize) + box = SecretBox(xfer_key) with open(filename, "rb") as f: plaintext = f.read() @@ -51,11 +54,14 @@ def send_file(so): transit_sender.add_their_relay_hints(tdata["relay_connection_hints"]) skt = transit_sender.establish_connection() - print("Sending %d bytes.." % filesize) + print("Sending..") + sent = 0 + next_update = start_progress(len(encrypted)) while sent < len(encrypted): - more = skt.send(encrypted[sent:]) - sent += more + sent += skt.send(encrypted[sent:]) + next_update = update_progress(next_update, sent, len(encrypted)) + finish_progress(len(encrypted)) print("File sent.. waiting for confirmation") # ack is a short newline-terminated string, followed by socket close. A long diff --git a/src/wormhole/scripts/progress.py b/src/wormhole/scripts/progress.py new file mode 100644 index 0000000..e74248d --- /dev/null +++ b/src/wormhole/scripts/progress.py @@ -0,0 +1,40 @@ +from __future__ import print_function +import sys, time + +def print_progress(completed, expected): + # scp does "<>(13% 168MB 39.3MB/s 00:27 ETA)" + # we do "Progress: #### 13% 168MB" + fmt = "Progress: %-40s %3d%% %4d%s" + short_unit_size, short_unit_name = 1, "B" + if expected > 9999: + short_unit_size, short_unit_name = 1000, "KB" + if expected > 9999*1000: + short_unit_size, short_unit_name = 1000*1000, "MB" + if expected > 9999*1000*1000: + short_unit_size, short_unit_name = 1000*1000*1000, "GB" + + percentage_complete = 1.0 * completed / expected + bars = "#" * int(percentage_complete * 40) + perc = int(100 * percentage_complete) + short_unit_count = int(completed / short_unit_size) + out = fmt % (bars, perc, short_unit_count, short_unit_name) + print("\r"+" "*70, end="") + print("\r"+out, end="") + sys.stdout.flush() + +def start_progress(expected, UPDATE_EVERY=0.2): + print_progress(0, expected) + next_update = time.time() + UPDATE_EVERY + return next_update + +def update_progress(next_update, completed, expected, UPDATE_EVERY=0.2): + now = time.time() + if now < next_update: + return next_update + next_update = now + UPDATE_EVERY + print_progress(completed, expected) + return next_update + +def finish_progress(expected): + print_progress(expected, expected) + print()