From 7699ed2291e8f73fc23410099b2f4058946c450b Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Wed, 5 Apr 2017 12:50:52 -0700
Subject: [PATCH] tell wordlist how many words to expect, add hyphens to
 matches

I'm still undecided about whether to add this to the mailbox
properties (revealing it to attackers) or continue to require non-default
wordcounts to be provided as a --code-length= argument to the receiver. So
for now the only place that says count=2 is in the default argument on
get_completions().
---
 src/wormhole/_wordlist.py          | 13 +++++++++----
 src/wormhole/test/test_wordlist.py | 13 +++++++++----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/wormhole/_wordlist.py b/src/wormhole/_wordlist.py
index ac48e78..cf801c5 100644
--- a/src/wormhole/_wordlist.py
+++ b/src/wormhole/_wordlist.py
@@ -1,4 +1,4 @@
-from __future__ import unicode_literals
+from __future__ import unicode_literals, print_function
 import os
 from zope.interface import implementer
 from ._interfaces import IWordlist
@@ -160,9 +160,10 @@ for k,both_words in raw_words.items():
 
 @implementer(IWordlist)
 class PGPWordList(object):
-    def get_completions(self, prefix):
+    def get_completions(self, prefix, num_words=2):
         # start with the odd words
-        if prefix.count("-") % 2 == 0:
+        count = prefix.count("-")
+        if count % 2 == 0:
             words = odd_words_lowercase
         else:
             words = even_words_lowercase
@@ -171,7 +172,11 @@ class PGPWordList(object):
         completions = set()
         for word in words:
             if word.startswith(last_partial_word):
-                completions.add(word[lp:])
+                suffix = word[lp:]
+                # append a hyphen if we expect more words
+                if count+1 < num_words:
+                    suffix += "-"
+                completions.add(suffix)
         return completions
 
     def choose_words(self, length):
diff --git a/src/wormhole/test/test_wordlist.py b/src/wormhole/test/test_wordlist.py
index 56cac4b..c86fbdb 100644
--- a/src/wormhole/test/test_wordlist.py
+++ b/src/wormhole/test/test_wordlist.py
@@ -7,11 +7,16 @@ class Completions(unittest.TestCase):
     def test_completions(self):
         wl = PGPWordList()
         gc = wl.get_completions
-        self.assertEqual(gc("ar"), {"mistice", "ticle"})
-        self.assertEqual(gc("armis"), {"tice"})
-        self.assertEqual(gc("armistice-ba"),
+        self.assertEqual(gc("ar", 2), {"mistice-", "ticle-"})
+        self.assertEqual(gc("armis", 2), {"tice-"})
+        self.assertEqual(gc("armistice", 2), {"-"})
+        self.assertEqual(gc("armistice-ba", 2),
                          {"boon", "ckfield", "ckward", "njo"})
-        self.assertEqual(gc("armistice-baboon"), {""})
+        self.assertEqual(gc("armistice-ba", 3),
+                         {"boon-", "ckfield-", "ckward-", "njo-"})
+        self.assertEqual(gc("armistice-baboon", 2), {""})
+        self.assertEqual(gc("armistice-baboon", 3), {"-"})
+        self.assertEqual(gc("armistice-baboon", 4), {"-"})
 
 class Choose(unittest.TestCase):
     def test_choose_words(self):