From db724a7b9158a6f5cfae658709125f8d30e21d3a Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Sun, 15 Feb 2015 09:16:17 -0800
Subject: [PATCH] simplify ask-user-for-code (with completion) API

---
 bin/receive_text.py        | 10 ++++------
 src/wormhole/transcribe.py | 35 ++++++++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/bin/receive_text.py b/bin/receive_text.py
index e4497c4..b8b25d6 100644
--- a/bin/receive_text.py
+++ b/bin/receive_text.py
@@ -1,17 +1,15 @@
 from __future__ import print_function
 import time, json
-from wormhole.transcribe import Receiver, list_channels
+from wormhole.transcribe import Receiver
 from wormhole.codes import input_code_with_completion
 
 APPID = "lothar.com/wormhole/text-xfer"
 
 # we're receiving
-channel_ids = list_channels()
-code = input_code_with_completion("Enter receive-text wormhole code: ",
-                                  channel_ids)
-start = time.time()
 data = json.dumps({"message": "ok"}).encode("utf-8")
-r = Receiver(APPID, data, code)
+r = Receiver(APPID, data)
+r.set_code(r.input_code("Enter receive-text wormhole code: "))
+start = time.time()
 them_bytes = r.get_data()
 them_d = json.loads(them_bytes.decode("utf-8"))
 print(them_d["message"])
diff --git a/src/wormhole/transcribe.py b/src/wormhole/transcribe.py
index dc8942c..d20a8d0 100644
--- a/src/wormhole/transcribe.py
+++ b/src/wormhole/transcribe.py
@@ -3,9 +3,9 @@ from binascii import hexlify, unhexlify
 from spake2 import SPAKE2_A, SPAKE2_B
 from nacl.secret import SecretBox
 from nacl import utils
+from . import codes
 from .hkdf import HKDF
 from .const import RELAY
-from .codes import make_code, extract_channel_id
 
 SECOND = 1
 MINUTE = 60*SECOND
@@ -95,7 +95,7 @@ class Initiator(Common):
 
     def get_code(self):
         self.channel_id = self._allocate() # allocate channel
-        self.code = make_code(self.channel_id)
+        self.code = codes.make_code(self.channel_id)
         self.sp = SPAKE2_A(self.code.encode("ascii"),
                            idA=self.appid+":Initiator",
                            idB=self.appid+":Receiver")
@@ -117,29 +117,42 @@ class Initiator(Common):
         return inbound_data
 
 
-def list_channels(relay=RELAY):
-    r = requests.get(relay + "list")
-    r.raise_for_status()
-    channel_ids = r.json()["channel-ids"]
-    return channel_ids
-
 class Receiver(Common):
-    def __init__(self, appid, data, code, relay=RELAY):
+    def __init__(self, appid, data, relay=RELAY):
         self.appid = appid
         self.data = data
-        self.code = code
-        self.channel_id = extract_channel_id(code)
         self.relay = relay
         assert relay.endswith("/")
         self.started = time.time()
         self.wait = 0.5*SECOND
         self.timeout = 3*MINUTE
         self.side = "receiver"
+        self.code = None
+        self.channel_id = None
+
+    def list_channels(self):
+        r = requests.get(self.relay + "list")
+        r.raise_for_status()
+        channel_ids = r.json()["channel-ids"]
+        return channel_ids
+
+    def input_code(self, prompt="Enter wormhole code: "):
+        channel_ids = self.list_channels()
+        code = codes.input_code_with_completion(prompt, channel_ids)
+        return code
+
+    def set_code(self, code):
+        assert self.code is None
+        assert self.channel_id is None
+        self.code = code
+        self.channel_id = codes.extract_channel_id(code)
         self.sp = SPAKE2_B(code.encode("ascii"),
                            idA=self.appid+":Initiator",
                            idB=self.appid+":Receiver")
 
     def get_data(self):
+        assert self.code is not None
+        assert self.channel_id is not None
         other_msgs = self._post_pake()
         key = self._poll_pake(other_msgs)