From 069b76485b3014edc6af21e4f3027c2bcc798567 Mon Sep 17 00:00:00 2001
From: meejah <meejah@meejah.ca>
Date: Sun, 14 Aug 2016 17:14:29 -0600
Subject: [PATCH] Add 'wormhole ssh-add' and 'wormhole ssh-send' commands

---
 src/wormhole/cli/cli.py     |  41 +++++++++++++++
 src/wormhole/cli/cmd_ssh.py |  72 +++++++++++++++++++++++++
 src/wormhole/xfer_util.py   | 101 ++++++++++++++++++++++++++++++++++++
 3 files changed, 214 insertions(+)
 create mode 100644 src/wormhole/cli/cmd_ssh.py
 create mode 100644 src/wormhole/xfer_util.py

diff --git a/src/wormhole/cli/cli.py b/src/wormhole/cli/cli.py
index 8e48bbd..ff706a0 100644
--- a/src/wormhole/cli/cli.py
+++ b/src/wormhole/cli/cli.py
@@ -30,6 +30,7 @@ class Config(object):
         self.cwd = os.getcwd()
         self.stdout = stdout
         self.stderr = stderr
+        self.tor = False  # XXX?
 
 def _compose(*decorators):
     def decorate(f):
@@ -217,3 +218,43 @@ def receive(cfg, code, **kwargs):
         cfg.code = None
 
     return go(cmd_receive.receive, cfg)
+
+
+@wormhole.command(name="ssh-add")
+@click.option(
+    "-c", "--code-length", default=2,
+    metavar="NUMWORDS",
+    help="length of code (in bytes/words)",
+)
+@click.option(
+    "--auth-file", "-f",
+    default=expanduser('~/.ssh/authorized_keys'),
+    type=click.Path(exists=False),
+)
+@click.pass_context
+def ssh_add(ctx, code_length, auth_file):
+    from . import cmd_ssh
+    ctx.obj.code_length = code_length
+    ctx.obj.auth_file = auth_file
+    return go(cmd_ssh.add, ctx.obj)
+
+
+@wormhole.command(name="ssh-send")
+@click.argument(
+    "code", nargs=1, required=True,
+)
+@click.option(
+    "--yes", "-y", is_flag=True,
+    help="Skip confirmation prompt to send key",
+)
+@click.pass_obj
+def ssh_send(cfg, code, yes):
+    from . import cmd_ssh
+    kind, keyid, pubkey = cmd_ssh.find_public_key()
+    print("Sending public key type='{}' keyid='{}'".format(kind, keyid))
+    if yes is not True:
+        click.confirm("Really send public key '{}' ?".format(keyid), abort=True)
+    cfg.public_key = (kind, keyid, pubkey)
+    cfg.code = code
+
+    return go(cmd_ssh.send, cfg)
diff --git a/src/wormhole/cli/cmd_ssh.py b/src/wormhole/cli/cmd_ssh.py
new file mode 100644
index 0000000..8b2faf3
--- /dev/null
+++ b/src/wormhole/cli/cmd_ssh.py
@@ -0,0 +1,72 @@
+from __future__ import print_function
+
+from os.path import expanduser, exists
+from twisted.internet.defer import inlineCallbacks
+from twisted.internet import reactor
+
+from .. import xfer_util
+
+
+def find_public_key():
+    """
+    This looks for an appropriate SSH key to send, possibly querying
+    the user in the meantime.
+
+    Returns a 3-tuple: kind, keyid, pubkey_data
+    """
+
+    # XXX FIXME don't blindly just send this one...
+    with open(expanduser('~/.ssh/id_rsa.pub'), 'r') as f:
+        pubkey = f.read()
+    parts = pubkey.strip().split()
+    kind = parts[0]
+    keyid = 'unknown' if len(parts) <= 2 else parts[2]
+
+    return kind, keyid, pubkey
+
+
+@inlineCallbacks
+def send(cfg, reactor=reactor):
+    yield xfer_util.send(
+        reactor,
+        u"lothar.com/wormhole/ssh-add",
+        cfg.relay_url,
+        data=cfg.public_key[2],
+        code=cfg.code,
+        use_tor=cfg.tor,
+    )
+    print("Key sent.")
+
+
+@inlineCallbacks
+def add(cfg, reactor=reactor):
+
+    def on_code_created(code):
+        print("Now tell the other user to run:")
+        print()
+        print("wormhole ssh-send {}".format(code))
+        print()
+
+    pubkey = yield xfer_util.receive(
+        reactor,
+        u"lothar.com/wormhole/ssh-add",
+        cfg.relay_url,
+        None,  # allocate a code for us
+        use_tor=cfg.tor,
+        on_code=on_code_created,
+    )
+
+    parts = pubkey.split()
+    kind = parts[0]
+    keyid = 'unknown' if len(parts) <= 2 else parts[2]
+
+    path = cfg.auth_file
+    if path == '-':
+        print(pubkey.strip())
+    else:
+        if not exists(path):
+            print("Note: '{}' not found; will be created".format(path))
+        with open(path, 'a') as f:
+            f.write('{}\n'.format(pubkey.strip()))
+        print("Appended key type='{kind}' id='{key_id}' to '{auth_file}'".format(
+            kind=kind, key_id=keyid, auth_file=path))
diff --git a/src/wormhole/xfer_util.py b/src/wormhole/xfer_util.py
new file mode 100644
index 0000000..41f9c9f
--- /dev/null
+++ b/src/wormhole/xfer_util.py
@@ -0,0 +1,101 @@
+import json
+from twisted.internet.defer import inlineCallbacks, returnValue
+
+from .wormhole import wormhole
+
+
+@inlineCallbacks
+def receive(reactor, appid, relay_url, code, use_tor=None, on_code=None):
+    """
+    This is a convenience API which returns a Deferred that callbacks
+    with a single chunk of data from another wormhole (and then closes
+    the wormhole). Under the hood, it's just using an instance
+    returned from :func:`wormhole.wormhole`. This is similar to the
+    `wormhole receive` command.
+
+    :param unicode appid: our application ID
+
+    :param unicode relay_url: the relay URL to use
+
+    :param unicode code: a pre-existing code to use, or None
+
+    :param bool use_tor: True if we should use Tor, False to not use it (None for default)
+
+    :param on_code: if not None, this is called when we have a code (even if you passed in one explicitly)
+    :type on_code: single-argument callable
+    """
+    wh = wormhole(appid, relay_url, reactor, use_tor)
+    if code is None:
+        code = yield wh.get_code()
+    else:
+        wh.set_code(code)
+    # we'll call this no matter what, even if you passed in a code --
+    # maybe it should be only in the 'if' block above?
+    if on_code:
+        on_code(code)
+    data = yield wh.get()
+    data = json.loads(data)
+    offer = data.get('offer', None)
+    if not offer:
+        raise Exception(
+            "Do not understand response: {}".format(data)
+        )
+    msg = None
+    if 'message' in offer:
+        msg = offer['message']
+        wh.send(json.dumps({"answer": {"message_ack": "ok"}}))
+
+    else:
+        raise Exception(
+            "Unknown offer type: {}".format(offer.keys())
+        )
+
+    yield wh.close()
+    returnValue(msg)
+
+
+@inlineCallbacks
+def send(reactor, appid, relay_url, data, code, use_tor=None, on_code=None):
+    """
+    This is a convenience API which returns a Deferred that callbacks
+    after a single chunk of data has been sent to another
+    wormhole. Under the hood, it's just using an instance returned
+    from :func:`wormhole.wormhole`. This is similar to the `wormhole
+    send` command.
+
+    :param unicode appid: the application ID
+
+    :param unicode relay_url: the relay URL to use
+
+    :param unicode code: a pre-existing code to use, or None
+
+    :param bool use_tor: True if we should use Tor, False to not use it (None for default)
+
+    :param on_code: if not None, this is called when we have a code (even if you passed in one explicitly)
+    :type on_code: single-argument callable
+    """
+    wh = wormhole(appid, relay_url, reactor, use_tor)
+    if code is None:
+        code = yield wh.get_code()
+    else:
+        wh.set_code(code)
+    if on_code:
+        on_code(code)
+
+    wh.send(
+        json.dumps({
+            "offer": {
+                "message": data
+            }
+        })
+    )
+    data = yield wh.get()
+    data = json.loads(data)
+    answer = data.get('answer', None)
+    yield wh.close()
+    if answer:
+        returnValue(None)
+    else:
+        raise Exception(
+            "Unknown answer: {}".format(data)
+        )