From 6fefcde0619b75b628f49f363c04cd8b90e0e91a Mon Sep 17 00:00:00 2001 From: Brian Warner Date: Wed, 7 Oct 2015 12:59:58 -0700 Subject: [PATCH] add send/receive of whole directories --- src/wormhole/scripts/cmd_receive.py | 88 ++++++++++++++++++++++++++++- src/wormhole/scripts/cmd_send.py | 58 +++++++++++++++---- src/wormhole/scripts/runner.py | 16 +++--- 3 files changed, 141 insertions(+), 21 deletions(-) diff --git a/src/wormhole/scripts/cmd_receive.py b/src/wormhole/scripts/cmd_receive.py index 10eab3d..79dbb24 100644 --- a/src/wormhole/scripts/cmd_receive.py +++ b/src/wormhole/scripts/cmd_receive.py @@ -1,5 +1,5 @@ from __future__ import print_function -import os, sys, json, binascii, six +import os, sys, json, binascii, six, tempfile, zipfile from ..errors import handle_server_error APPID = u"lothar.com/wormhole/text-or-file-xfer" @@ -77,6 +77,89 @@ def accept_file(args, them_d, w): record_pipe.close() return 0 +def accept_directory(args, them_d, w): + from ..blocking.transit import TransitReceiver, TransitError + from .progress import start_progress, update_progress, finish_progress + + file_data = them_d["directory"] + mode = file_data["mode"] + if mode != "zipfile/deflated": + print("Error: unknown directory-transfer mode '%s'" % (mode,)) + data = json.dumps({"error": "unknown mode"}).encode("utf-8") + w.send_data(data) + return 1 + + # the basename() is intended to protect us against + # "~/.ssh/authorized_keys" and other attacks + dirname = os.path.basename(file_data["dirname"]) # unicode + filesize = file_data["zipsize"] + num_files = file_data["numfiles"] + num_bytes = file_data["numbytes"] + + if os.path.exists(dirname): + print("Error: refusing to overwrite existing directory %s" % (dirname,)) + data = json.dumps({"error": "directory already exists"}).encode("utf-8") + w.send_data(data) + return 1 + + print("Receiving directory into: %s/" % (dirname,)) + print("%d files, %d bytes (%d compressed)" % (num_files, num_bytes, + filesize)) + while True and not args.accept_file: + ok = six.moves.input("ok? (y/n): ") + if ok.lower().startswith("y"): + break + print("transfer rejected", file=sys.stderr) + data = json.dumps({"error": "transfer rejected"}).encode("utf-8") + w.send_data(data) + return 1 + + transit_receiver = TransitReceiver(args.transit_helper) + data = json.dumps({ + "file_ack": "ok", + "transit": { + "direct_connection_hints": transit_receiver.get_direct_hints(), + "relay_connection_hints": transit_receiver.get_relay_hints(), + }, + }).encode("utf-8") + w.send_data(data) + # now done with the Wormhole object + + # now receive the rest of the owl + tdata = them_d["transit"] + transit_key = w.derive_key(APPID+u"/transit-key") + transit_receiver.set_transit_key(transit_key) + transit_receiver.add_their_direct_hints(tdata["direct_connection_hints"]) + transit_receiver.add_their_relay_hints(tdata["relay_connection_hints"]) + record_pipe = transit_receiver.connect() + + print("Receiving %d bytes for '%s' (%s).." % (filesize, dirname, + transit_receiver.describe())) + f = tempfile.SpooledTemporaryFile() + received = 0 + next_update = start_progress(filesize) + while received < filesize: + try: + plaintext = record_pipe.receive_record() + except TransitError: + print() + print("Connection dropped before full file received") + print("got %d bytes, wanted %d" % (received, filesize)) + return 1 + f.write(plaintext) + received += len(plaintext) + next_update = update_progress(next_update, received, filesize) + finish_progress(filesize) + assert received == filesize + print("Unpacking zipfile..") + with zipfile.ZipFile(f, "r", zipfile.ZIP_DEFLATED) as zf: + zf.extractall(path=dirname) + + print("Received files written to %s/" % dirname) + record_pipe.send_record(b"ok\n") + record_pipe.close() + return 0 + @handle_server_error def receive(args): # we're receiving text, or a file @@ -120,6 +203,9 @@ def receive(args): if "file" in them_d: return accept_file(args, them_d, w) + if "directory" in them_d: + return accept_directory(args, them_d, w) + print("I don't know what they're offering\n") print("Offer details:", them_d) data = json.dumps({"error": "unknown offer type"}).encode("utf-8") diff --git a/src/wormhole/scripts/cmd_send.py b/src/wormhole/scripts/cmd_send.py index 29bcc98..f026264 100644 --- a/src/wormhole/scripts/cmd_send.py +++ b/src/wormhole/scripts/cmd_send.py @@ -1,12 +1,12 @@ from __future__ import print_function -import os, sys, json, binascii, six +import os, sys, json, binascii, six, tempfile, zipfile from ..errors import handle_server_error APPID = u"lothar.com/wormhole/text-or-file-xfer" @handle_server_error def send(args): - # we're sending text, or a file + # we're sending text, or a file/directory from ..blocking.transcribe import Wormhole, WrongPasswordError from ..blocking.transit import TransitSender from .progress import start_progress, update_progress, finish_progress @@ -26,10 +26,9 @@ def send(args): "message": text, } else: - if not os.path.isfile(args.what): - print("Cannot send: no file named '%s'" % args.what) + if not os.path.exists(args.what): + print("Cannot send: no file/directory named '%s'" % args.what) return 1 - # we're sending a file sending_message = False transit_sender = TransitSender(args.transit_helper) phase1 = { @@ -38,14 +37,49 @@ def send(args): "relay_connection_hints": transit_sender.get_relay_hints(), }, } - filesize = os.stat(args.what).st_size basename = os.path.basename(args.what) - print("Sending %d byte file named '%s'" % (filesize, basename)) - fd_to_send = open(args.what, "rb") - phase1["file"] = { - "filename": basename, - "filesize": filesize, - } + if os.path.isfile(args.what): + # we're sending a file + filesize = os.stat(args.what).st_size + phase1["file"] = { + "filename": basename, + "filesize": filesize, + } + print("Sending %d byte file named '%s'" % (filesize, basename)) + fd_to_send = open(args.what, "rb") + elif os.path.isdir(args.what): + print("Building zipfile..") + # We're sending a directory. Create a zipfile in a tempdir and + # send that. + fd_to_send = tempfile.SpooledTemporaryFile() + # TODO: I think ZIP_DEFLATED means compressed.. check it + num_files = 0 + num_bytes = 0 + tostrip = len(args.what.split(os.sep)) + with zipfile.ZipFile(fd_to_send, "w", zipfile.ZIP_DEFLATED) as zf: + for path,dirs,files in os.walk(args.what): + # path always starts with args.what, then sometimes might + # have "/subdir" appended. We want the zipfile to contain + # "" or "subdir" + localpath = list(path.split(os.sep)[tostrip:]) + for fn in files: + archivename = os.path.join(*tuple(localpath+[fn])) + localfilename = os.path.join(path, fn) + zf.write(localfilename, archivename) + num_bytes += os.stat(localfilename).st_size + num_files += 1 + fd_to_send.seek(0,2) + filesize = fd_to_send.tell() + fd_to_send.seek(0,0) + phase1["directory"] = { + "mode": "zipfile/deflated", + "dirname": basename, + "zipsize": filesize, + "numbytes": num_bytes, + "numfiles": num_files, + } + print("Sending directory (%d bytes compressed) named '%s'" + % (filesize, basename)) with Wormhole(APPID, args.relay_url) as w: if args.zeromode: diff --git a/src/wormhole/scripts/runner.py b/src/wormhole/scripts/runner.py index d78205a..e4497b9 100644 --- a/src/wormhole/scripts/runner.py +++ b/src/wormhole/scripts/runner.py @@ -76,21 +76,21 @@ sp_tail_usage.set_defaults(func=cmd_usage.tail_usage) # CLI: send p = subparsers.add_parser("send", - description="Send text message or file", - usage="wormhole send [FILENAME]") + description="Send text message, file, or directory", + usage="wormhole send [FILENAME|DIRNAME]") p.add_argument("--text", metavar="MESSAGE", help="text message to send, instead of a file. Use '-' to read from stdin.") p.add_argument("--code", metavar="CODE", help="human-generated code phrase", type=type(u"")) p.add_argument("-0", dest="zeromode", action="store_true", help="enable no-code anything-goes mode") -p.add_argument("what", nargs="?", default=None, metavar="[FILENAME]", - help="the file to send") +p.add_argument("what", nargs="?", default=None, metavar="[FILENAME|DIRNAME]", + help="the file/directory to send") p.set_defaults(func=cmd_send.send) # CLI: receive p = subparsers.add_parser("receive", - description="Receive a text message or file", + description="Receive a text message, file, or directory", usage="wormhole receive [CODE]") p.add_argument("-0", dest="zeromode", action="store_true", help="enable no-code anything-goes mode") @@ -98,10 +98,10 @@ p.add_argument("-t", "--only-text", dest="only_text", action="store_true", help="refuse file transfers, only accept text transfers") p.add_argument("--accept-file", dest="accept_file", action="store_true", help="accept file transfer with asking for confirmation") -p.add_argument("-o", "--output-file", default=None, metavar="FILENAME", +p.add_argument("-o", "--output-file", default=None, metavar="FILENAME|DIRNAME", help=dedent("""\ - The file to create, overriding the filename suggested by the - sender."""), + The file or directory to create, overriding the name suggested + by the sender."""), ) p.add_argument("code", nargs="?", default=None, metavar="[CODE]", help=dedent("""\