diff --git a/src/wormhole/scripts/cmd_receive_twisted.py b/src/wormhole/scripts/cmd_receive_twisted.py new file mode 100644 index 0000000..166c10c --- /dev/null +++ b/src/wormhole/scripts/cmd_receive_twisted.py @@ -0,0 +1,154 @@ +from __future__ import print_function +import io, json +from zope.interface import implementer +from twisted.internet import interfaces, defer +from twisted.internet.defer import inlineCallbacks, returnValue +from ..twisted.transcribe import Wormhole, WrongPasswordError +from ..twisted.transit import TransitReceiver #, TransitError +from .cmd_receive_blocking import BlockingReceiver, RespondError, APPID +from ..errors import TransferError +from .progress import ProgressPrinter + +def receive_twisted(args): + return TwistedReceiver(args).go() + +class TwistedReceiver(BlockingReceiver): + + # TODO: @handle_server_error + @inlineCallbacks + def go(self): + w = Wormhole(APPID, self.args.relay_url) + + rc = yield self._go(w) + yield w.close() + returnValue(rc) + + @inlineCallbacks + def _go(self, w): + self.handle_code(w) + verifier = yield w.get_verifier() + self.show_verifier(verifier) + them_d = yield self.get_data(w) + try: + if "message" in them_d: + yield self.handle_text(them_d, w) + returnValue(0) + if "file" in them_d: + f = self.handle_file(them_d) + rp = yield self.establish_transit(w, them_d) + yield self.transfer_data(rp, f) + self.write_file(f) + yield self.close_transit(rp) + elif "directory" in them_d: + f = self.handle_directory(them_d) + rp = yield self.establish_transit(w, them_d) + yield self.transfer_data(rp, f) + self.write_directory(f) + yield self.close_transit(rp) + else: + self.msg(u"I don't know what they're offering\n") + self.msg(u"Offer details:", them_d) + raise RespondError({"error": "unknown offer type"}) + except RespondError as r: + data = json.dumps(r.response).encode("utf-8") + yield w.send_data(data) + returnValue(1) + returnValue(0) + + @inlineCallbacks + def get_data(self, w): + try: + them_bytes = yield w.get_data() + except WrongPasswordError as e: + raise TransferError(u"ERROR: " + e.explain()) + them_d = json.loads(them_bytes.decode("utf-8")) + if "error" in them_d: + raise TransferError(u"ERROR: " + them_d["error"]) + returnValue(them_d) + + @inlineCallbacks + def handle_text(self, them_d, w): + # we're receiving a text message + self.msg(them_d["message"]) + data = json.dumps({"message_ack": "ok"}).encode("utf-8") + yield w.send_data(data) + + @inlineCallbacks + def establish_transit(self, w, them_d): + transit_key = w.derive_key(APPID+u"/transit-key") + transit_receiver = TransitReceiver(self.args.transit_helper) + transit_receiver.set_transit_key(transit_key) + direct_hints = yield transit_receiver.get_direct_hints() + relay_hints = yield transit_receiver.get_relay_hints() + data = json.dumps({ + "file_ack": "ok", + "transit": { + "direct_connection_hints": direct_hints, + "relay_connection_hints": relay_hints, + }, + }).encode("utf-8") + yield w.send_data(data) + + # now receive the rest of the owl + tdata = them_d["transit"] + transit_receiver.add_their_direct_hints(tdata["direct_connection_hints"]) + transit_receiver.add_their_relay_hints(tdata["relay_connection_hints"]) + record_pipe = yield transit_receiver.connect() + returnValue(record_pipe) + + @inlineCallbacks + def transfer_data(self, record_pipe, f): + self.msg(u"Receiving (%s).." % record_pipe.describe()) + + progress_stdout = self.args.stdout + if self.args.hide_progress: + progress_stdout = io.StringIO() + pfc = ProgressingFileConsumer(f, self.xfersize, progress_stdout) + record_pipe.connectConsumer(pfc) + received = yield pfc.when_done + record_pipe.disconnectConsumer() + # except TransitError + if received < self.xfersize: + self.msg() + self.msg(u"Connection dropped before full file received") + self.msg(u"got %d bytes, wanted %d" % (received, self.xfersize)) + returnValue(1) # TODO: exit properly + assert received == self.xfersize + + @inlineCallbacks + def close_transit(self, record_pipe): + yield record_pipe.send_record(b"ok\n") + yield record_pipe.close() + +# based on twisted.protocols.ftp.FileConsumer, but: +# - finish after 'xfersize' bytes received, instead of connectionLost() +# - don't close the filehandle when done + +@implementer(interfaces.IConsumer) +class ProgressingFileConsumer: + def __init__(self, f, xfersize, progress_stdout): + self._f = f + self._xfersize = xfersize + self._received = 0 + self._progress = ProgressPrinter(xfersize, progress_stdout) + self._progress.start() + self.when_done = defer.Deferred() + + def registerProducer(self, producer, streaming): + self.producer = producer + assert streaming + + def write(self, bytes): + self._f.write(bytes) + self._received += len(bytes) + self._progress.update(self._received) + if self._received >= self._xfersize: + self._progress.finish() + d,self.when_done = self.when_done,None + d.callback(self._received) + + def unregisterProducer(self): + self.producer = None + if self.when_done: + # connection was dropped before all bytes were received + self.when_done.callback(self._received) diff --git a/src/wormhole/test/test_scripts.py b/src/wormhole/test/test_scripts.py index fe63b33..2d31dc7 100644 --- a/src/wormhole/test/test_scripts.py +++ b/src/wormhole/test/test_scripts.py @@ -8,7 +8,7 @@ from twisted.internet.threads import deferToThread from .. import __version__ from .common import ServerBase from ..scripts import (runner, cmd_send_blocking, cmd_send_twisted, - cmd_receive_blocking) + cmd_receive_blocking, cmd_receive_twisted) from ..scripts.send_common import build_phase1_data from ..errors import TransferError @@ -303,8 +303,10 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase): send_d = cmd_send_twisted.send_twisted(sargs) else: send_d = deferToThread(cmd_send_blocking.send_blocking, sargs) - assert not receiver_twisted # not importable yet - receive_d = deferToThread(cmd_receive_blocking.receive_blocking, rargs) + if receiver_twisted: + receive_d = cmd_receive_twisted.receive_twisted(rargs) + else: + receive_d = deferToThread(cmd_receive_blocking.receive_blocking, rargs) send_rc = yield send_d send_stdout = sargs.stdout.getvalue() @@ -391,6 +393,10 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase): return self._do_test(as_subprocess=True) def test_text_twisted_to_blocking(self): return self._do_test(sender_twisted=True) + def test_text_blocking_to_twisted(self): + return self._do_test(receiver_twisted=True) + def test_text_twisted_to_twisted(self): + return self._do_test(sender_twisted=True, receiver_twisted=True) def test_file(self): return self._do_test(mode="file") @@ -398,6 +404,11 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase): return self._do_test(mode="file", override_filename=True) def test_file_twisted_to_blocking(self): return self._do_test(mode="file", sender_twisted=True) + def test_file_blocking_to_twisted(self): + return self._do_test(mode="file", receiver_twisted=True) + def test_file_twisted_to_twisted(self): + return self._do_test(mode="file", + sender_twisted=True, receiver_twisted=True) def test_directory(self): return self._do_test(mode="directory") @@ -405,3 +416,8 @@ class PregeneratedCode(ServerBase, ScriptsBase, unittest.TestCase): return self._do_test(mode="directory", override_filename=True) def test_directory_twisted_to_blocking(self): return self._do_test(mode="directory", sender_twisted=True) + def test_directory_blocking_to_twisted(self): + return self._do_test(mode="directory", receiver_twisted=True) + def test_directory_twisted_to_twisted(self): + return self._do_test(mode="directory", + sender_twisted=True, receiver_twisted=True)