Island: Use IUserDatastore in AuthenticationService

This commit is contained in:
Mike Salvatore 2021-11-18 11:29:45 -05:00
parent a2eab2fa5d
commit b239b76063
3 changed files with 197 additions and 16 deletions

View File

@ -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
)

View File

@ -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))

View File

@ -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()