Skip to content
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# Generated by Django 4.1.5 on 2023-02-03 04:55

from django.db import migrations, models


class Migration(migrations.Migration):
dependencies = [
("testapp", "0003_remove_user_yubikey_id"),
]

operations = [
migrations.AlterField(
model_name="user",
name="first_name",
field=models.CharField(
blank=True, max_length=150, verbose_name="first name"
),
),
migrations.AlterField(
model_name="user",
name="id",
field=models.BigAutoField(
auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
),
),
migrations.AlterField(
model_name="user",
name="last_name",
field=models.CharField(
blank=True, max_length=150, verbose_name="last name"
),
),
]
15 changes: 15 additions & 0 deletions testproject/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ def mfa_method_creator(
is_primary=is_primary,
name=method_name,
is_active=method_args.pop("is_active", True),
code_generated_at=None,
**method_args,
)

Expand Down Expand Up @@ -74,6 +75,20 @@ def active_user_with_email_otp() -> UserModel:
return user


@pytest.fixture()
def active_user_with_email_hotp() -> UserModel:
user, created = User.objects.get_or_create(
username="imhotep",
email="imhotep@pyramids.eg",
)
if created:
user.set_password("secretkey"),
user.is_active = True
user.save()
mfa_method_creator(user=user, method_name="email", is_totp=False)
return user


@pytest.fixture()
def deactivated_user_with_email_otp(active_user_with_email_otp) -> UserModel:
active_user_with_email_otp.is_active = False
Expand Down
6 changes: 3 additions & 3 deletions testproject/tests/test_add_mfa.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from rest_framework.status import HTTP_200_OK, HTTP_400_BAD_REQUEST

from tests.utils import TrenchAPIClient
from trench.command.create_otp import create_otp_command
from trench.command.create_otp import create_totp_command
from trench.command.create_secret import create_secret_command


Expand All @@ -24,7 +24,7 @@ def test_add_user_mfa(active_user):
path="/auth/email/activate/",
data={
"secret": secret,
"code": create_otp_command(secret=secret, interval=60).now(),
"code": create_totp_command(secret=secret, interval=60).now(),
"user": getattr(active_user, active_user.USERNAME_FIELD),
},
format="json",
Expand All @@ -43,7 +43,7 @@ def test_should_fail_on_add_user_mfa_with_invalid_source_field(active_user: User
path="/auth/email/activate/",
data={
"secret": secret,
"code": create_otp_command(secret=secret, interval=60).now(),
"code": create_totp_command(secret=secret, interval=60).now(),
"user": getattr(active_user, active_user.USERNAME_FIELD),
},
format="json",
Expand Down
31 changes: 30 additions & 1 deletion testproject/tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,42 @@ def test_innermost_object_test(active_user):


@pytest.mark.django_db
def test_validate_code(active_user_with_email_otp):
def test_validate_code_totp(active_user_with_email_otp):
email_method = active_user_with_email_otp.mfa_methods.get()
handler = get_mfa_handler(mfa_method=email_method)
valid_code = handler.create_code()

assert handler.validate_code(code="123456") is False
assert handler.validate_code(code=valid_code) is True


@pytest.mark.django_db
def test_validate_code_hotp(active_user_with_email_hotp):
email_method = active_user_with_email_hotp.mfa_methods.get()
handler = get_mfa_handler(mfa_method=email_method)
valid_code = handler.create_code()

email_method.refresh_from_db()
assert email_method.code_generated_at is not None

assert handler.validate_code(code="123456") is False
email_method.refresh_from_db()
assert email_method.code_generated_at is not None

# successful validation clears code generation timestemp
assert handler.validate_code(code=valid_code) is True
email_method.refresh_from_db()
assert email_method.code_generated_at is None

# subsequently validating the same code twice fails
assert handler.validate_code(code=valid_code) is False

# creating a new code invalidates the previous one
valid_code = handler.create_code()
new_valid_code = handler.create_code()
assert new_valid_code != valid_code
assert handler.validate_code(code=valid_code) is False
assert handler.validate_code(code=new_valid_code) is True


@pytest.mark.django_db
Expand Down
48 changes: 41 additions & 7 deletions trench/backends/base.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from django.db.models import Model
from django.utils import timezone

from abc import ABC, abstractmethod
from pyotp import TOTP
from datetime import timedelta
from pyotp import TOTP, HOTP
from typing import Any, Dict, Optional, Tuple

from trench.command.create_otp import create_otp_command
from trench.command.create_otp import create_totp_command, create_hotp_command
from trench.exceptions import MissingConfigurationError
from trench.models import MFAMethod
from trench.responses import DispatchResponse
Expand Down Expand Up @@ -64,7 +66,15 @@ def dispatch_message(self) -> DispatchResponse:
raise NotImplementedError # pragma: no cover

def create_code(self) -> str:
return self._get_otp().now()
# totp
if self._mfa_method.is_totp:
return self._get_otp().now()

# hotp
Comment thread
nefrob marked this conversation as resolved.
Outdated
self._mfa_method.counter += 1
self._mfa_method.code_generated_at = timezone.now()
self._mfa_method.save()
return self._get_otp().at(self._mfa_method.counter)

def confirm_activation(self, code: str) -> None:
pass
Expand All @@ -73,12 +83,36 @@ def validate_confirmation_code(self, code: str) -> bool:
return self.validate_code(code)

def validate_code(self, code: str) -> bool:
return self._get_otp().verify(otp=code)
# totp
if self._mfa_method.is_totp:
return self._get_otp().verify(otp=code)

# hotp
is_valid = self._get_otp().verify(otp=code, counter=self._mfa_method.counter)
if not is_valid or not self._mfa_method.code_generated_at:
return False

def _get_otp(self) -> TOTP:
return create_otp_command(
secret=self._mfa_method.secret, interval=self._get_valid_window()
min_time = self._mfa_method.code_generated_at
max_time = self._mfa_method.code_generated_at + timedelta(
seconds=self._get_valid_window()
)
now = timezone.now()
if now < min_time or now > max_time:
return False

self._mfa_method.code_generated_at = None
self._mfa_method.save()
return True

def _get_otp(self) -> TOTP | HOTP:
# totp
if self._mfa_method.is_totp:
return create_totp_command(
secret=self._mfa_method.secret, interval=self._get_valid_window()
)

# hotp
return create_hotp_command(secret=self._mfa_method.secret)

def _get_valid_window(self) -> int:
return self._config.get(
Expand Down
3 changes: 2 additions & 1 deletion trench/command/create_mfa_method.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,11 @@ def __init__(self, secret_generator: Callable, mfa_model: Type[MFAMethod]) -> No
self._mfa_model = mfa_model
self._create_secret = secret_generator

def execute(self, user_id: int, name: str) -> MFAMethod:
def execute(self, user_id: int, name: str, is_totp: bool = True) -> MFAMethod:
mfa, created = self._mfa_model.objects.get_or_create(
user_id=user_id,
name=name,
is_totp=is_totp,
defaults={
"secret": self._create_secret,
"is_active": False,
Expand Down
15 changes: 12 additions & 3 deletions trench/command/create_otp.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,19 @@
from pyotp import TOTP
from pyotp import TOTP, HOTP


class CreateOTPCommand:
class CreateTOTPCommand:
@staticmethod
def execute(secret: str, interval: int) -> TOTP:
return TOTP(secret, interval=interval)


create_otp_command = CreateOTPCommand.execute
create_totp_command = CreateTOTPCommand.execute


class CreateHOTPCommand:
@staticmethod
def execute(secret: str) -> HOTP:
return HOTP(secret)


create_hotp_command = CreateHOTPCommand.execute
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
# Generated by Django 4.1.5 on 2023-02-03 04:55

from django.db import migrations, models


class Migration(migrations.Migration):
dependencies = [
("trench", "0004_alter_mfamethod_id_mfamethod_unique_user_is_primary_and_more"),
]

operations = [
migrations.RemoveConstraint(
model_name="mfamethod",
name="primary_is_active",
),
migrations.AddField(
model_name="mfamethod",
name="code_generated_at",
field=models.DateTimeField(
blank=True, null=True, verbose_name="code generated at"
),
),
migrations.AddField(
model_name="mfamethod",
name="counter",
field=models.PositiveIntegerField(default=0, verbose_name="counter"),
),
migrations.AddField(
model_name="mfamethod",
name="is_totp",
field=models.BooleanField(default=True, verbose_name="is totp"),
),
migrations.AlterField(
model_name="mfamethod",
name="id",
field=models.BigAutoField(
auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
),
),
migrations.AddConstraint(
model_name="mfamethod",
constraint=models.CheckConstraint(
check=models.Q(
models.Q(("is_primary", True), ("is_active", True)),
("is_primary", False),
_connector="OR",
),
name="primary_is_active",
),
),
]
7 changes: 6 additions & 1 deletion trench/models.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
from django.conf import settings
from django.db.models import (
CASCADE,
BooleanField,
CASCADE,
CharField,
CheckConstraint,
DateTimeField,
ForeignKey,
Manager,
Model,
PositiveIntegerField,
Q,
QuerySet,
TextField,
Expand Down Expand Up @@ -70,6 +72,9 @@ class MFAMethod(Model):
)
name = CharField(_("name"), max_length=255)
secret = CharField(_("secret"), max_length=255)
counter = PositiveIntegerField(_("counter"), default=0)
Comment thread
nefrob marked this conversation as resolved.
code_generated_at = DateTimeField(_("code generated at"), blank=True, null=True)
is_totp = BooleanField(_("is totp"), default=True)
is_primary = BooleanField(_("is primary"), default=False)
is_active = BooleanField(_("is active"), default=False)
_backup_codes = TextField(_("backup codes"), blank=True)
Expand Down
5 changes: 5 additions & 0 deletions trench/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def __getitem__(self, attr: str) -> Any:
SOURCE_FIELD = "SOURCE_FIELD"
HANDLER = "HANDLER"
VALIDITY_PERIOD = "VALIDITY_PERIOD"
USE_TOPT = "USE_TOPT"
VERBOSE_NAME = "VERBOSE_NAME"
EMAIL_SUBJECT = "EMAIL_SUBJECT"
EMAIL_PLAIN_TEMPLATE = "EMAIL_PLAIN_TEMPLATE"
Expand Down Expand Up @@ -73,13 +74,15 @@ def __getitem__(self, attr: str) -> Any:
"sms_twilio": {
VERBOSE_NAME: _("sms_twilio"),
VALIDITY_PERIOD: 30,
USE_TOPT: True,
HANDLER: "trench.backends.twilio.TwilioMessageDispatcher",
SOURCE_FIELD: "phone_number",
TWILIO_VERIFIED_FROM_NUMBER: "YOUR TWILIO REGISTERED NUMBER",
},
"sms_api": {
VERBOSE_NAME: _("sms_api"),
VALIDITY_PERIOD: 30,
USE_TOPT: True,
HANDLER: "trench.backends.sms_api.SMSAPIMessageDispatcher",
SOURCE_FIELD: "phone_number",
SMSAPI_ACCESS_TOKEN: "YOUR SMSAPI TOKEN",
Expand All @@ -88,6 +91,7 @@ def __getitem__(self, attr: str) -> Any:
"sms_aws": {
VERBOSE_NAME: _("sms_aws"),
VALIDITY_PERIOD: 30,
USE_TOPT: True,
HANDLER: "trench.backends.aws.AWSMessageDispatcher",
SOURCE_FIELD: "phone_number",
AWS_ACCESS_KEY: "YOUR AWS ACCESS KEY",
Expand All @@ -97,6 +101,7 @@ def __getitem__(self, attr: str) -> Any:
"email": {
VERBOSE_NAME: _("email"),
VALIDITY_PERIOD: 30,
USE_TOPT: True,
HANDLER: "trench.backends.basic_mail.SendMailMessageDispatcher",
SOURCE_FIELD: "email",
EMAIL_SUBJECT: _("Your verification code"),
Expand Down
10 changes: 6 additions & 4 deletions trench/views/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
MFAMethodDeactivationValidator,
UserMFAMethodSerializer,
)
from trench.settings import SOURCE_FIELD, trench_settings
from trench.settings import SOURCE_FIELD, USE_TOPT, trench_settings
from trench.utils import available_method_choices, get_mfa_model, user_token_generator


Expand Down Expand Up @@ -104,20 +104,22 @@ class MFAMethodActivationView(APIView):
@staticmethod
def post(request: Request, method: str) -> Response:
try:
source_field = get_mfa_config_by_name_query(name=method).get(SOURCE_FIELD)
mfa_config = get_mfa_config_by_name_query(name=method)
source_field = mfa_config.get(SOURCE_FIELD)
is_totp = mfa_config.get(USE_TOPT, True)
except MFAMethodDoesNotExistError as cause:
return ErrorResponse(error=cause)
user = request.user
try:
if source_field is not None and not hasattr(user, source_field):
raise MFASourceFieldDoesNotExistError(
source_field,
user.__class__.__name__
source_field, user.__class__.__name__
)

mfa = create_mfa_method_command(
user_id=user.id,
name=method,
is_totp=is_totp,
)
except MFAValidationError as cause:
return ErrorResponse(error=cause)
Expand Down