From 2f307423f0066ee01b83de2b69679edef0e8e069 Mon Sep 17 00:00:00 2001
From: Daniel Townsend <dan@dantownsend.co.uk>
Date: Mon, 9 Sep 2024 12:40:36 +0100
Subject: [PATCH] 295 Fix extra requirements for MFA (#296)

* fix extra requirements for MFA

* fix linter error
---
 piccolo_api/encryption/providers.py | 25 ++++++++++---------------
 setup.py                            |  7 ++++++-
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/piccolo_api/encryption/providers.py b/piccolo_api/encryption/providers.py
index fe1eaf5..f10073e 100644
--- a/piccolo_api/encryption/providers.py
+++ b/piccolo_api/encryption/providers.py
@@ -5,16 +5,16 @@
 from abc import ABCMeta, abstractmethod
 
 if t.TYPE_CHECKING:
-    import cryptography
     import nacl
+    from cryptography.fernet import Fernet
 
 
 logger = logging.getLogger(__name__)
 
 
-def get_cryptography() -> cryptography:  # type: ignore
+def get_fernet_class() -> t.Type[Fernet]:  # type: ignore
     try:
-        import cryptography
+        from cryptography.fernet import Fernet
     except ImportError as e:
         print(
             "Install pip install piccolo_api[cryptography] to use this "
@@ -22,7 +22,7 @@ def get_cryptography() -> cryptography:  # type: ignore
         )
         raise e
 
-    return cryptography
+    return Fernet
 
 
 class EncryptionProvider(metaclass=ABCMeta):
@@ -108,14 +108,12 @@ def __init__(self, encryption_key: bytes):
 
     @staticmethod
     def get_new_key() -> bytes:
-        cryptography = get_cryptography()
-        return cryptography.fernet.Fernet.generate_key()  # type: ignore
+        Fernet = get_fernet_class()
+        return Fernet.generate_key()  # type: ignore
 
     def encrypt(self, value: str, add_prefix: bool = True) -> str:
-        cryptography = get_cryptography()
-        fernet = cryptography.fernet.Fernet(  # type: ignore
-            self.encryption_key
-        )
+        Fernet = get_fernet_class()
+        fernet = Fernet(self.encryption_key)  # type: ignore
         encrypted_value = fernet.encrypt(value.encode("utf-8")).decode("utf-8")
         return (
             self.add_prefix(encrypted_value=encrypted_value)
@@ -127,11 +125,8 @@ def decrypt(self, encrypted_value: str, has_prefix: bool = True) -> str:
         if has_prefix:
             encrypted_value = self.remove_prefix(encrypted_value)
 
-        cryptography = get_cryptography()
-
-        fernet = cryptography.fernet.Fernet(  # type: ignore
-            self.encryption_key
-        )
+        Fernet = get_fernet_class()
+        fernet = Fernet(self.encryption_key)  # type: ignore
         return fernet.decrypt(encrypted_value.encode("utf-8")).decode("utf-8")
 
 
diff --git a/setup.py b/setup.py
index 8e416d2..d9abfe0 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,12 @@
     LONG_DESCRIPTION = f.read()
 
 
-EXTRAS = ["s3"]
+EXTRAS = [
+    "authenticator",
+    "cryptography",
+    "pynacl",
+    "s3",
+]
 
 
 def parse_requirement(req_path: str) -> t.List[str]: