def test_celery_beat(self, mock_expired_exceptions, mock_account_tech,
mock_purge, mock_setup):
from security_monkey.task_scheduler.beat import setup_the_tasks
from security_monkey.watchers.iam.iam_role import IAMRole
from security_monkey.auditors.iam.iam_role import IAMRoleAuditor
# Set up the monitor:
test_account = Account.query.filter(
Account.name == "TEST_ACCOUNT1").one()
watcher = IAMRole(accounts=[test_account.name])
batched_monitor = Monitor(IAMRole, test_account)
batched_monitor.watcher = watcher
batched_monitor.auditors = [
IAMRoleAuditor(accounts=[test_account.name])
]
import security_monkey.task_scheduler.tasks
old_get_monitors = security_monkey.task_scheduler.tasks.get_monitors
security_monkey.task_scheduler.tasks.get_monitors = lambda x, y, z: [
batched_monitor
]
setup_the_tasks(mock.Mock())
assert mock_setup.called
assert mock_purge.called
# "apply_async" where the immediately scheduled tasks
assert mock_account_tech.apply_async.called
# The ".s" are the scheduled tasks. Too lazy to grab the intervals out.
assert mock_account_tech.s.called
assert mock_expired_exceptions.s.called
#assert mock_task_audit.s.called
# Build the expected mock results:
scheduled_tech_result_list = []
async_result_list = []
# audit_result_list = []
import security_monkey.watcher
import security_monkey.auditor
for account in Account.query.filter(
Account.third_party == False).filter(
Account.active == True).all(): # noqa
for w in security_monkey.watcher.watcher_registry.iterkeys():
scheduled_tech_result_list.append(((account.name, w), ))
async_result_list.append((((account.name, w), ), ))
# It's just policy for IAM:
# audit_result_list.append(((account.name, "policy"),))
assert mock_account_tech.s.call_args_list == scheduled_tech_result_list
assert async_result_list == mock_account_tech.apply_async.call_args_list
# assert audit_result_list == mock_task_audit.s.call_args_list
security_monkey.task_scheduler.tasks.get_monitors = old_get_monitors
def test_celery_ignore_tech(self, mock_store_exception,
mock_expired_exceptions, mock_account_tech,
mock_purge, mock_setup):
import security_monkey.celeryconfig
security_monkey.celeryconfig.security_monkey_watcher_ignore = {
"policy"
}
from security_monkey.task_scheduler.beat import setup_the_tasks
from security_monkey.watchers.iam.iam_role import IAMRole
from security_monkey.watchers.iam.managed_policy import ManagedPolicy
from security_monkey.auditors.iam.iam_role import IAMRoleAuditor
from security_monkey.auditors.iam.iam_policy import IAMPolicyAuditor
# Stop the watcher registry from stepping on everyone's toes:
import security_monkey.watcher
import security_monkey.monitors
security_monkey.watcher.watcher_registry = {
IAMRole.index: IAMRole,
ManagedPolicy.index: ManagedPolicy
}
security_monkey.monitors.watcher_registry = security_monkey.watcher.watcher_registry
# Set up the monitors:
test_account = Account.query.filter(
Account.name == "TEST_ACCOUNT1").one()
role_watcher = IAMRole(accounts=[test_account.name])
mp_watcher = ManagedPolicy(accounts=[test_account.name])
batched_monitor = Monitor(IAMRole, test_account)
batched_monitor.watcher = role_watcher
batched_monitor.auditors = [
IAMRoleAuditor(accounts=[test_account.name])
]
normal_monitor = Monitor(ManagedPolicy, test_account)
normal_monitor.watcher = mp_watcher
normal_monitor.auditors = [
IAMPolicyAuditor(accounts=[test_account.name])
]
import security_monkey.task_scheduler.tasks
old_get_monitors = security_monkey.task_scheduler.tasks.get_monitors
security_monkey.task_scheduler.tasks.get_monitors = lambda x, y, z: [
batched_monitor, normal_monitor
]
setup_the_tasks(mock.Mock())
assert mock_setup.called
assert mock_purge.called
assert not mock_store_exception.called
# "apply_async" where the immediately scheduled tasks
assert mock_account_tech.apply_async.called
# The ".s" are the scheduled tasks. Too lazy to grab the intervals out.
assert mock_account_tech.s.called
assert mock_expired_exceptions.s.called
assert mock_expired_exceptions.apply_async.called
# Policy should not be called at all:
for mocked_call in mock_account_tech.s.call_args_list:
assert mocked_call[0][1] == "iamrole"
for mocked_call in mock_account_tech.apply_async.call_args_list:
assert mocked_call[0][0][1] == "iamrole"
security_monkey.task_scheduler.tasks.get_monitors = old_get_monitors
def test_celery_skipabeat(self, mock_store_exception,
mock_expired_exceptions, mock_account_tech,
mock_purge, mock_setup):
from security_monkey.task_scheduler.beat import setup_the_tasks
from security_monkey.watchers.github.org import GitHubOrg
from security_monkey.auditors.github.org import GitHubOrgAuditor
# Stop the watcher registry from stepping on everyone's toes:
import security_monkey.watcher
import security_monkey.monitors
security_monkey.watcher.watcher_registry = {GitHubOrg.index: GitHubOrg}
security_monkey.monitors.watcher_registry = security_monkey.watcher.watcher_registry
app.config["GITHUB_CREDENTIALS"] = {}
# Set up the monitor:
self.account_type = AccountType(name="GitHub")
db.session.add(self.account_type)
db.session.commit()
app.config["GITHUB_CREDENTIALS"] = {"Org-one": "token-one"}
db.session.add(
Account(name="Org-one",
account_type_id=self.account_type.id,
identifier="Org-one",
active=True,
third_party=False))
self.technology = Technology(name="organization")
db.session.add(self.technology)
db.session.commit()
# Disable the other accounts:
disable_account_1 = Account.query.filter(
Account.name == "TEST_ACCOUNT1").one()
disable_account_2 = Account.query.filter(
Account.name == "TEST_ACCOUNT2").one()
disable_account_1.active = False
disable_account_2.active = False
db.session.add(disable_account_1)
db.session.add(disable_account_2)
db.session.commit()
test_account = Account.query.filter(Account.name == "Org-one").one()
watcher = GitHubOrg(accounts=[test_account.name])
monitor = Monitor(GitHubOrg, test_account)
monitor.watcher = watcher
monitor.auditors = [GitHubOrgAuditor(accounts=[test_account.name])]
# This is externally executed (as in not with Celery):
db.session.add(
WatcherConfig(index=GitHubOrg.index, active=True, interval=0))
db.session.commit()
import security_monkey.task_scheduler.tasks
old_get_monitors = security_monkey.task_scheduler.tasks.get_monitors
security_monkey.task_scheduler.tasks.get_monitors = lambda x, y, z: [
monitor
]
get_interval = mock.Mock()
monitor.watcher.get_interval = get_interval
setup_the_tasks(mock.Mock())
assert mock_setup.called
assert mock_purge.called
assert not mock_store_exception.called
# "apply_async" will NOT be called...
assert not mock_account_tech.apply_async.called
# The ".s" are the scheduled tasks. Too lazy to grab the intervals out.
assert not mock_account_tech.s.called # Will not be called
assert mock_expired_exceptions.s.called
assert mock_expired_exceptions.apply_async.called
# Cleanup:
security_monkey.task_scheduler.tasks.get_monitors = old_get_monitors
disable_account_1.active = True
disable_account_2.active = True
test_account.active = False
db.session.add(disable_account_1)
db.session.add(disable_account_2)
db.session.add(test_account)
db.session.commit()