forked from p34709852/monkey
Island: Use IUserDatastore in AuthenticationService
This commit is contained in:
parent
a2eab2fa5d
commit
b239b76063
|
@ -1,7 +1,6 @@
|
|||
import bcrypt
|
||||
|
||||
import monkey_island.cc.environment.environment_singleton as env_singleton
|
||||
from common.utils.exceptions import IncorrectCredentialsError
|
||||
from common.utils.exceptions import IncorrectCredentialsError, UnknownUserError
|
||||
from monkey_island.cc.environment.user_creds import UserCreds
|
||||
from monkey_island.cc.server_utils.encryption import (
|
||||
reset_datastore_encryptor,
|
||||
|
@ -9,31 +8,40 @@ from monkey_island.cc.server_utils.encryption import (
|
|||
)
|
||||
from monkey_island.cc.setup.mongo.database_initializer import reset_database
|
||||
|
||||
from .i_user_datastore import IUserDatastore
|
||||
|
||||
|
||||
class AuthenticationService:
|
||||
DATA_DIR = None
|
||||
user_datastore = None
|
||||
|
||||
# TODO: A number of these services should be instance objects instead of
|
||||
# static/singleton hybrids. At the moment, this requires invasive refactoring that's
|
||||
# not a priority.
|
||||
@classmethod
|
||||
def initialize(cls, data_dir: str):
|
||||
def initialize(cls, data_dir: str, user_datastore: IUserDatastore):
|
||||
cls.DATA_DIR = data_dir
|
||||
cls.user_datastore = user_datastore
|
||||
|
||||
@staticmethod
|
||||
def needs_registration() -> bool:
|
||||
return env_singleton.env.needs_registration()
|
||||
@classmethod
|
||||
def needs_registration(cls) -> bool:
|
||||
return cls.user_datastore.has_registered_users()
|
||||
|
||||
@classmethod
|
||||
def register_new_user(cls, username: str, password: str):
|
||||
credentials = UserCreds(username, _hash_password(password))
|
||||
env_singleton.env.try_add_user(credentials)
|
||||
cls.user_datastore.add_user(credentials)
|
||||
cls._reset_datastore_encryptor(username, password)
|
||||
reset_database()
|
||||
|
||||
@classmethod
|
||||
def authenticate(cls, username: str, password: str):
|
||||
if not _credentials_match_registered_user(username, password):
|
||||
try:
|
||||
registered_user = cls.user_datastore.get_user_credentials(username)
|
||||
except UnknownUserError:
|
||||
raise IncorrectCredentialsError()
|
||||
|
||||
if not _credentials_match_registered_user(username, password, registered_user):
|
||||
raise IncorrectCredentialsError()
|
||||
|
||||
cls._unlock_datastore_encryptor(username, password)
|
||||
|
@ -56,12 +64,9 @@ def _hash_password(plaintext_password: str) -> str:
|
|||
return password_hash.decode()
|
||||
|
||||
|
||||
def _credentials_match_registered_user(username: str, password: str) -> bool:
|
||||
registered_user = env_singleton.env.get_user()
|
||||
|
||||
if not registered_user:
|
||||
return False
|
||||
|
||||
def _credentials_match_registered_user(
|
||||
username: str, password: str, registered_user: UserCreds
|
||||
) -> bool:
|
||||
return (registered_user.username == username) and _password_matches_hash(
|
||||
password, registered_user.password_hash
|
||||
)
|
||||
|
|
|
@ -1,10 +1,10 @@
|
|||
from monkey_island.cc.services.post_breach_files import PostBreachFilesService
|
||||
from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
|
||||
|
||||
from . import AuthenticationService
|
||||
from . import AuthenticationService, JsonFileUserDatastore
|
||||
|
||||
|
||||
def initialize_services(data_dir):
|
||||
PostBreachFilesService.initialize(data_dir)
|
||||
LocalMonkeyRunService.initialize(data_dir)
|
||||
AuthenticationService.initialize(data_dir)
|
||||
AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir))
|
||||
|
|
|
@ -0,0 +1,176 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from common.utils.exceptions import (
|
||||
AlreadyRegisteredError,
|
||||
IncorrectCredentialsError,
|
||||
InvalidRegistrationCredentialsError,
|
||||
UnknownUserError,
|
||||
)
|
||||
from monkey_island.cc.environment.user_creds import UserCreds
|
||||
from monkey_island.cc.services import AuthenticationService
|
||||
from monkey_island.cc.services.authentication import authentication_service
|
||||
from monkey_island.cc.services.authentication.i_user_datastore import IUserDatastore
|
||||
|
||||
USERNAME = "user1"
|
||||
PASSWORD = "test"
|
||||
PASSWORD_HASH = "$2b$12$YsGjjuJFddYJ6z5S5/nMCuKkCzKHB1AWY9SXkQ02i25d8TgdhIRS2"
|
||||
|
||||
|
||||
class MockUserDatastore(IUserDatastore):
|
||||
def __init__(self, has_registered_users, add_user, get_user_credentials):
|
||||
self._has_registered_users = has_registered_users
|
||||
self._add_user = add_user
|
||||
self._get_user_credentials = get_user_credentials
|
||||
|
||||
def has_registered_users(self):
|
||||
return self._has_registered_users()
|
||||
|
||||
def add_user(self, credentials: UserCreds):
|
||||
return self._add_user(credentials)
|
||||
|
||||
def get_user_credentials(self, username: str) -> UserCreds:
|
||||
return self._get_user_credentials(username)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_reset_datastore_encryptor():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_reset_database():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_unlock_datastore_encryptor():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_datastore_utils(
|
||||
monkeypatch,
|
||||
mock_reset_datastore_encryptor,
|
||||
mock_reset_database,
|
||||
mock_unlock_datastore_encryptor,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
authentication_service, "reset_datastore_encryptor", mock_reset_datastore_encryptor
|
||||
)
|
||||
monkeypatch.setattr(authentication_service, "reset_database", mock_reset_database)
|
||||
monkeypatch.setattr(
|
||||
authentication_service, "unlock_datastore_encryptor", mock_unlock_datastore_encryptor
|
||||
)
|
||||
|
||||
|
||||
# pass a mock IUserDatastore
|
||||
|
||||
# Mock reset_database
|
||||
# mock reset_datastore_encryptor
|
||||
# mock unlock_datastore_encryptor
|
||||
|
||||
|
||||
def test_needs_registration__true(tmp_path):
|
||||
mock_user_datastore = MockUserDatastore(lambda: True, None, None)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
assert a_s.needs_registration()
|
||||
|
||||
|
||||
def test_needs_registration__false(tmp_path):
|
||||
mock_user_datastore = MockUserDatastore(lambda: False, None, None)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
assert not a_s.needs_registration()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error", [InvalidRegistrationCredentialsError, AlreadyRegisteredError])
|
||||
def test_register_new_user__fails(
|
||||
tmp_path, mock_reset_datastore_encryptor, mock_reset_database, error
|
||||
):
|
||||
mock_user_datastore = MockUserDatastore(lambda: True, MagicMock(side_effect=error), None)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
with pytest.raises(error):
|
||||
a_s.register_new_user(USERNAME, PASSWORD)
|
||||
|
||||
mock_reset_datastore_encryptor.assert_not_called()
|
||||
mock_reset_database.assert_not_called()
|
||||
|
||||
|
||||
def test_register_new_user(tmp_path, mock_reset_datastore_encryptor, mock_reset_database):
|
||||
mock_add_user = MagicMock()
|
||||
mock_user_datastore = MockUserDatastore(lambda: False, mock_add_user, None)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
a_s.register_new_user(USERNAME, PASSWORD)
|
||||
|
||||
assert mock_add_user.call_args[0][0].username == USERNAME
|
||||
assert mock_add_user.call_args[0][0].password_hash != PASSWORD
|
||||
|
||||
mock_reset_datastore_encryptor.assert_called_once()
|
||||
assert mock_reset_datastore_encryptor.call_args[0][1] != USERNAME
|
||||
|
||||
mock_reset_database.assert_called_once()
|
||||
|
||||
|
||||
def test_authenticate__success(tmp_path, mock_unlock_datastore_encryptor):
|
||||
mock_user_datastore = MockUserDatastore(
|
||||
lambda: True,
|
||||
None,
|
||||
lambda _: UserCreds(USERNAME, PASSWORD_HASH),
|
||||
)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
# If authentication fails, this function will raise an exception and the test will fail.
|
||||
a_s.authenticate(USERNAME, PASSWORD)
|
||||
mock_unlock_datastore_encryptor.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("username", "password"), [("wrong_username", PASSWORD), (USERNAME, "wrong_password")]
|
||||
)
|
||||
def test_authenticate__failed_wrong_credentials(
|
||||
tmp_path, mock_unlock_datastore_encryptor, username, password
|
||||
):
|
||||
mock_user_datastore = MockUserDatastore(
|
||||
lambda: True,
|
||||
None,
|
||||
lambda _: UserCreds(USERNAME, PASSWORD_HASH),
|
||||
)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
# If authentication fails, this function will raise an exception and the test will fail.
|
||||
with pytest.raises(IncorrectCredentialsError):
|
||||
a_s.authenticate(username, password)
|
||||
|
||||
mock_unlock_datastore_encryptor.assert_not_called()
|
||||
|
||||
|
||||
def test_authenticate__failed_no_registered_user(tmp_path, mock_unlock_datastore_encryptor):
|
||||
mock_user_datastore = MockUserDatastore(
|
||||
lambda: True, None, MagicMock(side_effect=UnknownUserError)
|
||||
)
|
||||
|
||||
a_s = AuthenticationService()
|
||||
a_s.initialize(tmp_path, mock_user_datastore)
|
||||
|
||||
# If authentication fails, this function will raise an exception and the test will fail.
|
||||
with pytest.raises(IncorrectCredentialsError):
|
||||
a_s.authenticate(USERNAME, PASSWORD)
|
||||
|
||||
mock_unlock_datastore_encryptor.assert_not_called()
|
Loading…
Reference in New Issue