From 23c61fb92b7bb1aa56174de964ba6f41ee0c28ad Mon Sep 17 00:00:00 2001
From: Yorick Downe <yorickdowne@gmail.com>
Date: Thu, 19 Nov 2020 18:57:57 -0500
Subject: [PATCH] Allow 4-character abbreviations of mnemonic words when using
 `existing-mnemonic`

---
 eth2deposit/cli/existing_mnemonic.py          |  3 ++-
 .../key_handling/key_derivation/mnemonic.py   | 26 ++++++++++++-------
 ..._menmonic.py => test_existing_mnemonic.py} |  0
 .../test_key_derivation/test_mnemonic.py      |  8 +++---
 4 files changed, 22 insertions(+), 15 deletions(-)
 rename tests/test_cli/{test_existing_menmonic.py => test_existing_mnemonic.py} (100%)

diff --git a/eth2deposit/cli/existing_mnemonic.py b/eth2deposit/cli/existing_mnemonic.py
index 2c0970c2..c3ec741d 100644
--- a/eth2deposit/cli/existing_mnemonic.py
+++ b/eth2deposit/cli/existing_mnemonic.py
@@ -17,7 +17,8 @@
 
 
 def validate_mnemonic(cts: click.Context, param: Any, mnemonic: str) -> str:
-    if verify_mnemonic(mnemonic, WORD_LISTS_PATH):
+    mnemonic = verify_mnemonic(mnemonic, WORD_LISTS_PATH)
+    if mnemonic is not None:
         return mnemonic
     else:
         raise ValidationError('That is not a valid mnemonic, please check for typos.')
diff --git a/eth2deposit/key_handling/key_derivation/mnemonic.py b/eth2deposit/key_handling/key_derivation/mnemonic.py
index 756ec5eb..8ece460e 100644
--- a/eth2deposit/key_handling/key_derivation/mnemonic.py
+++ b/eth2deposit/key_handling/key_derivation/mnemonic.py
@@ -86,9 +86,10 @@ def determine_mnemonic_language(mnemonic: str, words_path: str) -> Sequence[str]
     There are collisions between word-lists, so multiple candidate languages are returned.
     """
     languages = get_languages(words_path)
-    word_language_map = {word: lang for lang in languages for word in _get_word_list(lang, words_path)}
+    word_language_map = {normalize('NFKC', word)[:4]: lang for lang in languages
+                         for word in _get_word_list(lang, words_path)}
     try:
-        mnemonic_list = mnemonic.split(' ')
+        mnemonic_list = [normalize('NFKC', word)[:4] for word in mnemonic.lower().split(' ')]
         word_languages = [word_language_map[word] for word in mnemonic_list]
         return list(set(word_languages))
     except KeyError:
@@ -110,30 +111,35 @@ def _get_checksum(entropy: bytes) -> int:
     return int.from_bytes(SHA256(entropy), 'big') >> (256 - checksum_length)
 
 
-def verify_mnemonic(mnemonic: str, words_path: str) -> bool:
+def verify_mnemonic(mnemonic: str, words_path: str) -> str:
     """
-    Given a mnemonic, verify it against its own checksum."
+    Given a mnemonic, verify it against its own checksum and return
+    a reconstructed full version - useful in case it was abbreviated.
     """
     try:
         languages = determine_mnemonic_language(mnemonic, words_path)
     except ValueError:
-        return False
+        return None
     for language in languages:
         try:
-            word_list = _get_word_list(language, words_path)
-            mnemonic_list = mnemonic.split(' ')
+            word_list = [normalize('NFKC', word)[:4] for word in _get_word_list(language, words_path)]
+            mnemonic_list = [normalize('NFKC', word)[:4] for word in mnemonic.lower().split(' ')]
             if len(mnemonic_list) not in range(12, 25, 3):
-                return False
+                return None
             word_indices = [_word_to_index(word_list, word) for word in mnemonic_list]
             mnemonic_int = _uint11_array_to_uint(word_indices)
             checksum_length = len(mnemonic_list) // 3
             checksum = mnemonic_int & 2**checksum_length - 1
             entropy = (mnemonic_int - checksum) >> checksum_length
             entropy_bits = entropy.to_bytes(checksum_length * 4, 'big')
-            return _get_checksum(entropy_bits) == checksum
+            full_word_list = _get_word_list(language, words_path)
+            if _get_checksum(entropy_bits) == checksum:
+                return ' '.join([_index_to_word(full_word_list, index) for index in word_indices])
+            else:
+                return None
         except ValueError:
             pass
-    return False
+    return None
 
 
 def get_mnemonic(*, language: str, words_path: str, entropy: Optional[bytes]=None) -> str:
diff --git a/tests/test_cli/test_existing_menmonic.py b/tests/test_cli/test_existing_mnemonic.py
similarity index 100%
rename from tests/test_cli/test_existing_menmonic.py
rename to tests/test_cli/test_existing_mnemonic.py
diff --git a/tests/test_key_handling/test_key_derivation/test_mnemonic.py b/tests/test_key_handling/test_key_derivation/test_mnemonic.py
index 51f6d335..0efa1580 100644
--- a/tests/test_key_handling/test_key_derivation/test_mnemonic.py
+++ b/tests/test_key_handling/test_key_derivation/test_mnemonic.py
@@ -38,13 +38,13 @@ def test_bip39(language: str, test: Sequence[str]) -> None:
 
 
 @pytest.mark.parametrize(
-    'test_mnemonic,is_valid',
-    [(test_mnemonic[1], True)
+    'test_mnemonic',
+    [(test_mnemonic[1])
      for _, language_test_vectors in test_vectors.items()
      for test_mnemonic in language_test_vectors]
 )
-def test_verify_mnemonic(test_mnemonic: str, is_valid: bool) -> None:
-    assert verify_mnemonic(test_mnemonic, WORD_LISTS_PATH) == is_valid
+def test_verify_mnemonic(test_mnemonic: str) -> None:
+    assert verify_mnemonic(test_mnemonic, WORD_LISTS_PATH) is not None
 
 
 @pytest.mark.parametrize(