add method allowing you to find the keys in a document

client/utils/encryption.py

@@ -5,14 +5,18 @@
 """
 import base64
 import os
+import re
 
 import pyaes
 
 HEADER_TEXT = "OKPY ENCRYPTED FILE FOLLOWS\n" + "-" * 100 + "\n"
 
-# used to ensure that the key us correct (helps detect incorrect key usage)
+# used to ensure that the key is correct (helps detect incorrect key usage)
 PLAINTEXT_PADDING = b"0" * 16
 
+# matches keys
+KEY_PATTERN = r"[a-z2-7]{52}9999"
+
 
 def generate_key() -> str:
     """
@@ -21,6 +25,20 @@ def generate_key() -> str:
     return to_safe_string(os.urandom(32))
 
 
+def is_valid_key(key: str) -> bool:
+    """
+    Returns if this is a valid key
+    """
+    return re.match("^" + KEY_PATTERN + "$", key) is not None
+
+
+def get_keys(document: str) -> list:
+    """
+    Gets all valid keys in the given document
+    """
+    return re.findall(KEY_PATTERN, document)
+
+
 def encrypt(data: str, key: str) -> str:
     """
     Encrypt the given data using the given key. Tag the result so that it is clear that this is an encrypted file.

tests/utils/encryption_test.py

@@ -36,3 +36,18 @@ def invalid_key_test(self):
     def key_characters_test(self):
         for _ in range(100):
             self.assertRegex(encryption.generate_key(), "^[0-9a-z]+$")
+            self.assertTrue(encryption.is_valid_key(encryption.generate_key()))
+
+    def find_single_key_test(self):
+        key = encryption.generate_key()
+        self.assertEqual([key], encryption.get_keys(
+            "some text some text some text {} some text some text some text".format(key)))
+        self.assertEqual([key], encryption.get_keys(key))
+
+        self.assertEqual([key] * 7, encryption.get_keys(key * 7))
+
+    def find_multiple_key_test(self):
+        key_a, key_b, key_c = [encryption.generate_key() for _ in range(3)]
+        self.assertEqual([key_a, key_b, key_c], encryption.get_keys(
+            "Key A: {}, Key B: {}, Key C: {}".format(key_a, key_b, key_c)))
+        self.assertEqual([key_a, key_c, key_b, key_a], encryption.get_keys(key_a + key_c + key_b + key_a))