add send/receive of whole directories

This commit is contained in:
Brian Warner 2015-10-07 12:59:58 -07:00
parent 6958241f3f
commit 6fefcde061
3 changed files with 141 additions and 21 deletions

View File

@ -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")

View File

@ -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:

View File

@ -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("""\