Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

3078 otp generator #3187

Merged
merged 10 commits into from
Apr 4, 2023
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,6 @@
from .flask_resources import register_resources

from .setup import setup_authentication

from .i_otp_generator import IOTPGenerator
from .authentication_otp_generator import AuthenticationOTPGenerator
Original file line number Diff line number Diff line change
@@ -1,17 +1,23 @@
import string
import time
from typing import Tuple

from flask_security import UserDatastore

from common.utils.code_utils import secure_generate_random_string
from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
from monkey_island.cc.models import IslandMode
from monkey_island.cc.server_utils.encryption import ILockableEncryptor
from monkey_island.cc.services.authentication_service.token_generator import TokenGenerator

from . import AccountRole
from .i_otp_repository import IOTPRepository
from .token_parser import ParsedToken, TokenParser
from .types import Token
from .types import OTP, Token
from .user import User

OTP_EXPIRATION_TIME = 2 * 60 # 2 minutes


class AuthenticationFacade:
"""
Expand All @@ -25,12 +31,14 @@ def __init__(
user_datastore: UserDatastore,
token_generator: TokenGenerator,
token_parser: TokenParser,
otp_repository: IOTPRepository,
):
self._repository_encryptor = repository_encryptor
self._island_event_queue = island_event_queue
self._datastore = user_datastore
self._token_generator = token_generator
self._token_parser = token_parser
self._otp_repository = otp_repository

def needs_registration(self) -> bool:
"""
Expand Down Expand Up @@ -71,6 +79,18 @@ def _get_refresh_token_owner(self, refresh_token: ParsedToken) -> User:
raise Exception("Invalid refresh token")
return user

def generate_otp(self) -> OTP:
"""
Generates a new OTP

The generated OTP is saved to the `IOTPRepository`
"""
otp = secure_generate_random_string(32, string.ascii_letters + string.digits + "._-")
expiration_time = time.monotonic() + OTP_EXPIRATION_TIME
self._otp_repository.insert_otp(otp, expiration_time)

return otp

def generate_refresh_token(self, user: User) -> Token:
"""
Generates a refresh token for a specific user
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from .authentication_facade import AuthenticationFacade
from .i_otp_generator import IOTPGenerator
from .types import OTP


class AuthenticationOTPGenerator(IOTPGenerator):
"""
Generates OTPs
"""

def __init__(self, authentication_facade: AuthenticationFacade):
self._authentication_facade = authentication_facade

def generate_otp(self) -> OTP:
"""
Generates a new OTP
"""
return self._authentication_facade.generate_otp()
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
from abc import ABC, abstractmethod

from .types import OTP


class IOTPGenerator(ABC):
"""Generator for OTPs"""

@abstractmethod
def generate_otp(self) -> OTP:
"""
Generates a new OTP
"""
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from unittest.mock import MagicMock

from monkey_island.cc.services.authentication_service import AuthenticationOTPGenerator
from monkey_island.cc.services.authentication_service.authentication_facade import (
AuthenticationFacade,
)


def test_authentication_otp_generator__generates_otp():
mock_authentication_facade = MagicMock(spec=AuthenticationFacade)

otp_generator = AuthenticationOTPGenerator(mock_authentication_facade)
otp_generator.generate_otp()

assert mock_authentication_facade.generate_otp.called_once
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import time
from unittest.mock import MagicMock, call

import pytest
Expand All @@ -8,8 +9,10 @@
from monkey_island.cc.models import IslandMode
from monkey_island.cc.server_utils.encryption import ILockableEncryptor
from monkey_island.cc.services.authentication_service.authentication_facade import (
OTP_EXPIRATION_TIME,
AuthenticationFacade,
)
from monkey_island.cc.services.authentication_service.i_otp_repository import IOTPRepository
from monkey_island.cc.services.authentication_service.setup import setup_authentication
from monkey_island.cc.services.authentication_service.token_generator import TokenGenerator
from monkey_island.cc.services.authentication_service.token_parser import (
Expand Down Expand Up @@ -58,6 +61,11 @@ def mock_token_parser() -> TokenParser:
return MagicMock(spec=TokenParser)


@pytest.fixture
def mock_otp_repository() -> IOTPRepository:
return MagicMock(spec=IOTPRepository)


@pytest.fixture
def authentication_facade(
mock_flask_app,
Expand All @@ -66,13 +74,15 @@ def authentication_facade(
mock_user_datastore: UserDatastore,
mock_token_generator: TokenGenerator,
mock_token_parser: TokenParser,
mock_otp_repository: IOTPRepository,
) -> AuthenticationFacade:
return AuthenticationFacade(
mock_repository_encryptor,
mock_island_event_queue,
mock_user_datastore,
mock_token_generator,
mock_token_parser,
mock_otp_repository,
)


Expand Down Expand Up @@ -178,6 +188,24 @@ def test_revoke_all_tokens_for_all_users(
[mock_user_datastore.set_uniquifier.assert_any_call(user) for user in USERS]


def test_generate_otp__saves_otp(
authentication_facade: AuthenticationFacade, mock_otp_repository: IOTPRepository
):
otp = authentication_facade.generate_otp()

assert mock_otp_repository.insert_otp.called_once_with(otp)


def test_generate_otp__uses_expected_expiration_time(
freezer, authentication_facade: AuthenticationFacade, mock_otp_repository: IOTPRepository
):
authentication_facade.generate_otp()

expiration_time = mock_otp_repository.insert_otp.call_args[0][1]
expected_expiration_time = time.monotonic() + OTP_EXPIRATION_TIME
assert expiration_time == expected_expiration_time


def test_setup_authentication__revokes_tokens(
mock_island_event_queue: IIslandEventQueue,
mock_repository_encryptor: ILockableEncryptor,
Expand Down
2 changes: 2 additions & 0 deletions vulture_allowlist.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
MongoAgentEventRepository,
MongoOTPRepository,
)
from monkey_island.cc.services.authentication_service import AuthenticationOTPGenerator
from monkey_island.cc.services.authentication_service.token import TokenValidator
from monkey_island.cc.services.authentication_service.user import User
from monkey_island.cc.services.reporting.exploitations.monkey_exploitation import MonkeyExploitation
Expand Down Expand Up @@ -160,3 +161,4 @@
IOTPRepository.get_expiration
IOTPRepository.reset
MongoOTPRepository
AuthenticationOTPGenerator.generate_otp