cmd_receive_twisted.py: twisted-based implementation

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)

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)