p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Refactored try_add_user and needs_registration to avoid code duplication

This commit is contained in:
VakarisZ 2020-06-22 11:59:02 +03:00
parent 8c428aa44d
commit 78cf0b5791
6 changed files with 149 additions and 129 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
monkey/common/utils/exceptions.py
View File

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

42
monkey/monkey_island/cc/environment/__init__.py
View File

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

113
monkey/monkey_island/cc/environment/test__init__.py
View File

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

64
monkey/monkey_island/cc/environment/test_environment_config.py
View File

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

4
monkey/monkey_island/cc/resources/registration.py
View File

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

41
monkey/monkey_island/cc/testing/environment/server_config_mocks.py Normal file
View File

@ -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"
}