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().
This commit is contained in:
Brian Warner 2017-04-05 12:50:52 -07:00
parent a1f0d1bbf7
commit 7699ed2291
2 changed files with 18 additions and 8 deletions

View File

@ -1,4 +1,4 @@
from __future__ import unicode_literals from __future__ import unicode_literals, print_function
import os import os
from zope.interface import implementer from zope.interface import implementer
from ._interfaces import IWordlist from ._interfaces import IWordlist
@ -160,9 +160,10 @@ for k,both_words in raw_words.items():
@implementer(IWordlist) @implementer(IWordlist)
class PGPWordList(object): class PGPWordList(object):
def get_completions(self, prefix): def get_completions(self, prefix, num_words=2):
# start with the odd words # start with the odd words
if prefix.count("-") % 2 == 0: count = prefix.count("-")
if count % 2 == 0:
words = odd_words_lowercase words = odd_words_lowercase
else: else:
words = even_words_lowercase words = even_words_lowercase
@ -171,7 +172,11 @@ class PGPWordList(object):
completions = set() completions = set()
for word in words: for word in words:
if word.startswith(last_partial_word): 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 return completions
def choose_words(self, length): def choose_words(self, length):

View File

@ -7,11 +7,16 @@ class Completions(unittest.TestCase):
def test_completions(self): def test_completions(self):
wl = PGPWordList() wl = PGPWordList()
gc = wl.get_completions gc = wl.get_completions
self.assertEqual(gc("ar"), {"mistice", "ticle"}) self.assertEqual(gc("ar", 2), {"mistice-", "ticle-"})
self.assertEqual(gc("armis"), {"tice"}) self.assertEqual(gc("armis", 2), {"tice-"})
self.assertEqual(gc("armistice-ba"), self.assertEqual(gc("armistice", 2), {"-"})
self.assertEqual(gc("armistice-ba", 2),
{"boon", "ckfield", "ckward", "njo"}) {"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): class Choose(unittest.TestCase):
def test_choose_words(self): def test_choose_words(self):