From 78cf0b57915d2fed6c8ebdb3c3b7431f60a4f3d5 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 22 Jun 2020 11:59:02 +0300 Subject: [PATCH] Refactored try_add_user and needs_registration to avoid code duplication --- monkey/common/utils/exceptions.py | 14 ++- .../monkey_island/cc/environment/__init__.py | 42 ++++--- .../cc/environment/test__init__.py | 113 +++++++++--------- .../cc/environment/test_environment_config.py | 64 ++-------- .../cc/resources/registration.py | 4 +- .../environment/server_config_mocks.py | 41 +++++++ 6 files changed, 149 insertions(+), 129 deletions(-) create mode 100644 monkey/monkey_island/cc/testing/environment/server_config_mocks.py diff --git a/monkey/common/utils/exceptions.py b/monkey/common/utils/exceptions.py index 5be984dfc..fa026933c 100644 --- a/monkey/common/utils/exceptions.py +++ b/monkey/common/utils/exceptions.py @@ -6,5 +6,17 @@ class FailedExploitationError(Exception): """ Raise when exploiter fails instead of returning False """ -class InvalidRegistrationCredentials(Exception): +class InvalidRegistrationCredentialsError(Exception): """ Raise when server config file changed and island needs to restart """ + + +class RegistrationNotNeededError(Exception): + """ Raise to indicate the reason why registration is not required """ + + +class CredentialsNotRequiredError(RegistrationNotNeededError): + """ Raise to indicate the reason why registration is not required """ + + +class AlreadyRegisteredError(RegistrationNotNeededError): + """ Raise to indicate the reason why registration is not required """ diff --git a/monkey/monkey_island/cc/environment/__init__.py b/monkey/monkey_island/cc/environment/__init__.py index 0e05c4923..5680b40e6 100644 --- a/monkey/monkey_island/cc/environment/__init__.py +++ b/monkey/monkey_island/cc/environment/__init__.py @@ -1,15 +1,18 @@ import hashlib +import logging import os from abc import ABCMeta, abstractmethod from datetime import timedelta __author__ = 'itay.mizeretz' - -from common.utils.exceptions import InvalidRegistrationCredentials +from common.utils.exceptions import InvalidRegistrationCredentialsError, RegistrationNotNeededError, \ + CredentialsNotRequiredError, AlreadyRegisteredError from monkey_island.cc.environment.environment_config import EnvironmentConfig from monkey_island.cc.environment.user_creds import UserCreds +logger = logging.getLogger(__name__) + class Environment(object, metaclass=ABCMeta): _ISLAND_PORT = 5000 @@ -36,24 +39,29 @@ class Environment(object, metaclass=ABCMeta): def get_auth_users(self): pass - def try_add_user(self, credentials: UserCreds): - if self._credentials_required: - if credentials: - if self._is_registered(): - raise InvalidRegistrationCredentials("User has already been registered. " - "Reset credentials or login.") - self._config.add_user(credentials) - else: - raise InvalidRegistrationCredentials("Missing part of credentials.") - else: - raise InvalidRegistrationCredentials("Can't add user because credentials are not required " - "for current environment.") - def needs_registration(self) -> bool: - if not self._credentials_required: + try: + needs_registration = self._try_needs_registration() + return needs_registration + except (CredentialsNotRequiredError, AlreadyRegisteredError) as e: + logger.info(e) return False + + def try_add_user(self, credentials: UserCreds): + if not credentials: + raise InvalidRegistrationCredentialsError("Missing part of credentials.") + if self._try_needs_registration(): + self._config.add_user(credentials) + + def _try_needs_registration(self) -> bool: + if not self._credentials_required: + raise CredentialsNotRequiredError("Credentials are not required " + "for current environment.") else: - return not self._is_registered() + if self._is_registered(): + raise AlreadyRegisteredError("User has already been registered. " + "Reset credentials or login.") + return True def _is_registered(self) -> bool: return self._credentials_required and self._is_credentials_set_up() diff --git a/monkey/monkey_island/cc/environment/test__init__.py b/monkey/monkey_island/cc/environment/test__init__.py index 3115023a5..870673013 100644 --- a/monkey/monkey_island/cc/environment/test__init__.py +++ b/monkey/monkey_island/cc/environment/test__init__.py @@ -4,8 +4,10 @@ from typing import Dict from unittest import TestCase from unittest.mock import patch, MagicMock -from common.utils.exceptions import InvalidRegistrationCredentials +from common.utils.exceptions import InvalidRegistrationCredentialsError, AlreadyRegisteredError, \ + CredentialsNotRequiredError, RegistrationNotNeededError from monkey_island.cc.environment import Environment, EnvironmentConfig, UserCreds +from monkey_island.cc.testing.environment.server_config_mocks import * def get_server_config_file_path_test_version(): @@ -14,94 +16,91 @@ def get_server_config_file_path_test_version(): class TestEnvironment(TestCase): - class EnvironmentNoCredentials(Environment): + class EnvironmentCredentialsNotRequired(Environment): + def __init__(self): + config = EnvironmentConfig('test', 'test', UserCreds()) + super().__init__(config) + _credentials_required = False def get_auth_users(self): return [] - class EnvironmentWithCredentials(Environment): + class EnvironmentCredentialsRequired(Environment): + def __init__(self): + config = EnvironmentConfig('test', 'test', UserCreds()) + super().__init__(config) + _credentials_required = True def get_auth_users(self): return [] - # Username:test Password:test - CONFIG_WITH_CREDENTIALS = { - "server_config": "password", - "deployment": "develop", - "user": "test", - "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" - "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" - } + class EnvironmentAlreadyRegistered(Environment): + def __init__(self): + config = EnvironmentConfig('test', 'test', UserCreds('test_user', 'test_secret')) + super().__init__(config) - CONFIG_NO_CREDENTIALS = { - "server_config": "password", - "deployment": "develop" - } + _credentials_required = True - CONFIG_PARTIAL_CREDENTIALS = { - "server_config": "password", - "deployment": "develop", - "user": "test" - } - - CONFIG_STANDARD_ENV = { - "server_config": "standard", - "deployment": "develop" - } - - CONFIG_STANDARD_WITH_CREDENTIALS = { - "server_config": "standard", - "deployment": "develop", - "user": "test", - "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" - "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" - } + def get_auth_users(self): + return [1, "Test_username", "Test_secret"] @patch.object(target=EnvironmentConfig, attribute="save_to_file", new=MagicMock()) def test_try_add_user(self): - env = TestEnvironment.EnvironmentWithCredentials() + env = TestEnvironment.EnvironmentCredentialsRequired() credentials = UserCreds(username="test", password_hash="1231234") env.try_add_user(credentials) credentials = UserCreds(username="test") - with self.assertRaises(InvalidRegistrationCredentials): + with self.assertRaises(InvalidRegistrationCredentialsError): env.try_add_user(credentials) - env = TestEnvironment.EnvironmentNoCredentials() + env = TestEnvironment.EnvironmentCredentialsNotRequired() credentials = UserCreds(username="test", password_hash="1231234") - with self.assertRaises(InvalidRegistrationCredentials): + with self.assertRaises(RegistrationNotNeededError): env.try_add_user(credentials) + def test_try_needs_registration(self): + env = TestEnvironment.EnvironmentAlreadyRegistered() + with self.assertRaises(AlreadyRegisteredError): + env._try_needs_registration() + + env = TestEnvironment.EnvironmentCredentialsNotRequired() + with self.assertRaises(CredentialsNotRequiredError): + env._try_needs_registration() + + env = TestEnvironment.EnvironmentCredentialsRequired() + self.assertTrue(env._try_needs_registration()) + def test_needs_registration(self): - env = TestEnvironment.EnvironmentWithCredentials() - self._test_bool_env_method("needs_registration", env, TestEnvironment.CONFIG_WITH_CREDENTIALS, False) - self._test_bool_env_method("needs_registration", env, TestEnvironment.CONFIG_NO_CREDENTIALS, True) - self._test_bool_env_method("needs_registration", env, TestEnvironment.CONFIG_PARTIAL_CREDENTIALS, True) + env = TestEnvironment.EnvironmentCredentialsRequired() + self._test_bool_env_method("needs_registration", env, CONFIG_WITH_CREDENTIALS, False) + self._test_bool_env_method("needs_registration", env, CONFIG_NO_CREDENTIALS, True) + self._test_bool_env_method("needs_registration", env, CONFIG_PARTIAL_CREDENTIALS, True) - env = TestEnvironment.EnvironmentNoCredentials() - self._test_bool_env_method("needs_registration", env, TestEnvironment.CONFIG_STANDARD_ENV, False) - self._test_bool_env_method("needs_registration", env, TestEnvironment.CONFIG_STANDARD_WITH_CREDENTIALS, False) + env = TestEnvironment.EnvironmentCredentialsNotRequired() + self._test_bool_env_method("needs_registration", env, CONFIG_STANDARD_ENV, False) + self._test_bool_env_method("needs_registration", env, CONFIG_STANDARD_WITH_CREDENTIALS, False) def test_is_registered(self): - env = TestEnvironment.EnvironmentWithCredentials() - self._test_bool_env_method("_is_registered", env, TestEnvironment.CONFIG_WITH_CREDENTIALS, True) - self._test_bool_env_method("_is_registered", env, TestEnvironment.CONFIG_NO_CREDENTIALS, False) - self._test_bool_env_method("_is_registered", env, TestEnvironment.CONFIG_PARTIAL_CREDENTIALS, False) + env = TestEnvironment.EnvironmentCredentialsRequired() + self._test_bool_env_method("_is_registered", env, CONFIG_WITH_CREDENTIALS, True) + self._test_bool_env_method("_is_registered", env, CONFIG_NO_CREDENTIALS, False) + self._test_bool_env_method("_is_registered", env, CONFIG_PARTIAL_CREDENTIALS, False) - env = TestEnvironment.EnvironmentNoCredentials() - self._test_bool_env_method("_is_registered", env, TestEnvironment.CONFIG_STANDARD_ENV, False) - self._test_bool_env_method("_is_registered", env, TestEnvironment.CONFIG_STANDARD_WITH_CREDENTIALS, False) + env = TestEnvironment.EnvironmentCredentialsNotRequired() + self._test_bool_env_method("_is_registered", env, CONFIG_STANDARD_ENV, False) + self._test_bool_env_method("_is_registered", env, CONFIG_STANDARD_WITH_CREDENTIALS, False) def test_is_credentials_set_up(self): - env = TestEnvironment.EnvironmentWithCredentials() - self._test_bool_env_method("_is_credentials_set_up", env, TestEnvironment.CONFIG_NO_CREDENTIALS, False) - self._test_bool_env_method("_is_credentials_set_up", env, TestEnvironment.CONFIG_WITH_CREDENTIALS, True) - self._test_bool_env_method("_is_credentials_set_up", env, TestEnvironment.CONFIG_PARTIAL_CREDENTIALS, False) + env = TestEnvironment.EnvironmentCredentialsRequired() + self._test_bool_env_method("_is_credentials_set_up", env, CONFIG_NO_CREDENTIALS, False) + self._test_bool_env_method("_is_credentials_set_up", env, CONFIG_WITH_CREDENTIALS, True) + self._test_bool_env_method("_is_credentials_set_up", env, CONFIG_PARTIAL_CREDENTIALS, False) - env = TestEnvironment.EnvironmentNoCredentials() - self._test_bool_env_method("_is_credentials_set_up", env, TestEnvironment.CONFIG_STANDARD_ENV, False) + env = TestEnvironment.EnvironmentCredentialsNotRequired() + self._test_bool_env_method("_is_credentials_set_up", env, CONFIG_STANDARD_ENV, False) def _test_bool_env_method(self, method_name: str, env: Environment, config: Dict, expected_result: bool): env._config = EnvironmentConfig.get_from_json(json.dumps(config)) diff --git a/monkey/monkey_island/cc/environment/test_environment_config.py b/monkey/monkey_island/cc/environment/test_environment_config.py index 2def7c4d3..2a4dbb2bb 100644 --- a/monkey/monkey_island/cc/environment/test_environment_config.py +++ b/monkey/monkey_island/cc/environment/test_environment_config.py @@ -8,6 +8,7 @@ from unittest.mock import patch, MagicMock from monkey_island.cc.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.environment.environment_config import EnvironmentConfig from monkey_island.cc.environment.user_creds import UserCreds +from monkey_island.cc.testing.environment.server_config_mocks import * def get_server_config_file_path_test_version(): @@ -15,52 +16,11 @@ def get_server_config_file_path_test_version(): class TestEnvironmentConfig(TestCase): - # Username:test Password:test - CONFIG_WITH_CREDENTIALS = { - "server_config": "password", - "deployment": "develop", - "user": "test", - "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" - "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" - } - - CONFIG_NO_CREDENTIALS = { - "server_config": "password", - "deployment": "develop" - } - - CONFIG_PARTIAL_CREDENTIALS = { - "server_config": "password", - "deployment": "develop", - "user": "test" - } - - CONFIG_BOGUS_VALUES = { - "server_config": "password", - "deployment": "develop", - "user": "test", - "aws": "test", - "test": "test", - "test2": "test2" - } - - CONFIG_STANDARD_ENV = { - "server_config": "standard", - "deployment": "develop" - } - - CONFIG_STANDARD_WITH_CREDENTIALS = { - "server_config": "standard", - "deployment": "develop", - "user": "test", - "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" - "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" - } def test_get_from_json(self): - self._test_get_from_json(TestEnvironmentConfig.CONFIG_WITH_CREDENTIALS) - self._test_get_from_json(TestEnvironmentConfig.CONFIG_NO_CREDENTIALS) - self._test_get_from_json(TestEnvironmentConfig.CONFIG_PARTIAL_CREDENTIALS) + self._test_get_from_json(CONFIG_WITH_CREDENTIALS) + self._test_get_from_json(CONFIG_NO_CREDENTIALS) + self._test_get_from_json(CONFIG_PARTIAL_CREDENTIALS) def _test_get_from_json(self, config: Dict): config_json = json.dumps(config) @@ -75,9 +35,9 @@ class TestEnvironmentConfig(TestCase): self.assertEqual(config['aws'], env_config_object.aws) def test_save_to_file(self): - self._test_save_to_file(TestEnvironmentConfig.CONFIG_WITH_CREDENTIALS) - self._test_save_to_file(TestEnvironmentConfig.CONFIG_NO_CREDENTIALS) - self._test_save_to_file(TestEnvironmentConfig.CONFIG_PARTIAL_CREDENTIALS) + self._test_save_to_file(CONFIG_WITH_CREDENTIALS) + self._test_save_to_file(CONFIG_NO_CREDENTIALS) + self._test_save_to_file(CONFIG_PARTIAL_CREDENTIALS) @patch.object(target=EnvironmentConfig, attribute="get_config_file_path", new=MagicMock(return_value=get_server_config_file_path_test_version())) @@ -103,14 +63,14 @@ class TestEnvironmentConfig(TestCase): self.assertEqual(EnvironmentConfig.get_config_file_path(), server_file_path) def test_get_from_dict(self): - config_dict = TestEnvironmentConfig.CONFIG_WITH_CREDENTIALS + config_dict = CONFIG_WITH_CREDENTIALS env_conf = EnvironmentConfig.get_from_dict(config_dict) self.assertEqual(env_conf.server_config, config_dict['server_config']) self.assertEqual(env_conf.deployment, config_dict['deployment']) self.assertEqual(env_conf.user_creds.username, config_dict['user']) self.assertEqual(env_conf.aws, None) - config_dict = TestEnvironmentConfig.CONFIG_BOGUS_VALUES + config_dict = CONFIG_BOGUS_VALUES env_conf = EnvironmentConfig.get_from_dict(config_dict) self.assertEqual(env_conf.server_config, config_dict['server_config']) self.assertEqual(env_conf.deployment, config_dict['deployment']) @@ -118,13 +78,13 @@ class TestEnvironmentConfig(TestCase): self.assertEqual(env_conf.aws, config_dict['aws']) def test_to_dict(self): - conf_json1 = json.dumps(TestEnvironmentConfig.CONFIG_WITH_CREDENTIALS) + conf_json1 = json.dumps(CONFIG_WITH_CREDENTIALS) self._test_to_dict(EnvironmentConfig.get_from_json(conf_json1)) - conf_json2 = json.dumps(TestEnvironmentConfig.CONFIG_NO_CREDENTIALS) + conf_json2 = json.dumps(CONFIG_NO_CREDENTIALS) self._test_to_dict(EnvironmentConfig.get_from_json(conf_json2)) - conf_json3 = json.dumps(TestEnvironmentConfig.CONFIG_PARTIAL_CREDENTIALS) + conf_json3 = json.dumps(CONFIG_PARTIAL_CREDENTIALS) self._test_to_dict(EnvironmentConfig.get_from_json(conf_json3)) def _test_to_dict(self, env_config_object: EnvironmentConfig): diff --git a/monkey/monkey_island/cc/resources/registration.py b/monkey/monkey_island/cc/resources/registration.py index 45351ccd2..46d50d2d0 100644 --- a/monkey/monkey_island/cc/resources/registration.py +++ b/monkey/monkey_island/cc/resources/registration.py @@ -1,7 +1,7 @@ import flask_restful from flask import request, make_response -from common.utils.exceptions import InvalidRegistrationCredentials +from common.utils.exceptions import InvalidRegistrationCredentialsError, RegistrationNotNeededError import monkey_island.cc.environment.environment_singleton as env_singleton from monkey_island.cc.environment.user_creds import UserCreds @@ -15,6 +15,6 @@ class Registration(flask_restful.Resource): try: env_singleton.env.try_add_user(credentials) return make_response({"error": ""}, 200) - except InvalidRegistrationCredentials as e: + except (InvalidRegistrationCredentialsError, RegistrationNotNeededError) as e: return make_response({"error": str(e)}, 400) diff --git a/monkey/monkey_island/cc/testing/environment/server_config_mocks.py b/monkey/monkey_island/cc/testing/environment/server_config_mocks.py new file mode 100644 index 000000000..ddbff3118 --- /dev/null +++ b/monkey/monkey_island/cc/testing/environment/server_config_mocks.py @@ -0,0 +1,41 @@ +# Username:test Password:test +CONFIG_WITH_CREDENTIALS = { + "server_config": "password", + "deployment": "develop", + "user": "test", + "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" + "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" +} + +CONFIG_NO_CREDENTIALS = { + "server_config": "password", + "deployment": "develop" +} + +CONFIG_PARTIAL_CREDENTIALS = { + "server_config": "password", + "deployment": "develop", + "user": "test" +} + +CONFIG_BOGUS_VALUES = { + "server_config": "password", + "deployment": "develop", + "user": "test", + "aws": "test", + "test": "test", + "test2": "test2" +} + +CONFIG_STANDARD_ENV = { + "server_config": "standard", + "deployment": "develop" +} + +CONFIG_STANDARD_WITH_CREDENTIALS = { + "server_config": "standard", + "deployment": "develop", + "user": "test", + "password_hash": "9ece086e9bac491fac5c1d1046ca11d737b92a2b2ebd93f005d7b710110c0a678288166e7fbe796883a" + "4f2e9b3ca9f484f521d0ce464345cc1aec96779149c14" +}