Verification revision

sanic_security/models.py

@@ -2,6 +2,7 @@
 import datetime
 import logging
 import re
+import uuid
 from typing import Union
 
 import jwt
@@ -20,7 +21,6 @@
     get_expiration_date,
     image_generator,
     audio_generator,
-    get_id,
     is_expired,
 )
 
@@ -58,7 +58,9 @@ class BaseModel(Model):
         deleted (bool): Renders the model filterable without removing from the database.
     """
 
-    id: str = fields.CharField(pk=True, max_length=36, default=get_id)
+    id: str = fields.CharField(
+        pk=True, max_length=36, default=lambda: str(uuid.uuid4())
+    )
     date_created: datetime.datetime = fields.DatetimeField(auto_now_add=True)
     date_updated: datetime.datetime = fields.DatetimeField(auto_now=True)
     deleted: bool = fields.BooleanField(default=False)
@@ -540,6 +542,10 @@ class Meta:
 class TwoStepSession(VerificationSession):
     """Validates a client using a code sent via email or text."""
 
+    code: str = fields.CharField(
+        max_length=6, default=lambda: get_code(True), null=True
+    )
+
     @classmethod
     async def new(cls, request: Request, account: Account, **kwargs):
         return await cls.create(

sanic_security/utils.py

@@ -1,7 +1,6 @@
 import datetime
 import random
-import string
-import uuid
+from string import ascii_uppercase, digits
 
 from argon2 import PasswordHasher
 from captcha.audio import AudioCaptcha
@@ -53,28 +52,22 @@ def get_ip(request: Request) -> str:
     return request.remote_addr or request.ip
 
 
-def get_code() -> str:
+def get_code(digits_only: bool = False) -> str:
     """
     Generates random code to be used for verification.
 
+    Args:
+        digits_only: Determines if code should only contain digits.
+
     Returns:
         code
     """
     return "".join(
-        random.choice(string.ascii_uppercase + string.digits) for _ in range(6)
+        random.choice(("" if digits_only else ascii_uppercase) + digits)
+        for _ in range(6)
     )
 
 
-def get_id() -> str:
-    """
-    Generates uuid to be used for primary key.
-
-    Returns:
-        id
-    """
-    return str(uuid.uuid4())
-
-
 def is_expired(date):
     """
     Checks if current date has surpassed the date passed into the function.