def test_unload_policy_file_with_bad_schema(self, os_unlink, _1, _2, _3,
_4, _5):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# Schedule is invalid
policy_string = """
policie:
- name: foo
mode:
type: container-periodic
schedule: '* * * *'
resource: azure.resourcegroup
"""
# Create a bad yaml file
file_path = tempfile.mktemp(suffix=".yaml")
with open(file_path, 'w') as f:
f.write(policy_string)
host.unload_policy_file(file_path, {})
os_unlink.assert_called()
# Clean up the file
os.remove(file_path)
def test_update_policies_ignore_policy_if_failed_to_create_content_hash(
self, get_blob_client_mock, _1, _2, _3, _4):
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", None), # no hash
]
client_mock.get_blob_to_path = self.download_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
def create_blob_from_bytes(_1, _2, _3, **kwargs):
raise AzureHttpError("Failed to create blob", 403)
client_mock.create_blob_from_bytes = create_blob_from_bytes
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# cleanup
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
self.assertEqual({}, host.policies)
# run
host.update_policies()
# no policies were loaded
self.assertEqual(0, len(host.policies.items()))
# jobs were created
jobs = host.scheduler.get_jobs()
self.assertEqual(0, len([j for j in jobs if j.id == 'blob1.yml']))
def test_update_policies(self, get_blob_client_mock, _1, _2, _3, _4):
# mock blob list call
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", "hash1"),
ContainerHostTest.get_mock_blob("blob2.YAML", "hash2"),
ContainerHostTest.get_mock_blob("blob3.md", "hash3")
]
client_mock.get_blob_to_path = self.download_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
# init
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# cleanup
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
self.assertEqual({}, host.policies)
# run
host.update_policies()
# both policies were loaded
self.assertEqual(2, len(host.policies.items()))
# jobs were created
jobs = host.scheduler.get_jobs()
self.assertEqual(1, len([j for j in jobs if j.id == 'blob1.yml']))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob2.YAML']))
def test_update_policies_create_content_hash(self, get_blob_client_mock,
_1, _2, _3, _4):
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", None), # no hash
]
client_mock.get_blob_to_path = self.download_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
def get_blob_properties(_, blob_name):
return ContainerHostTest.get_mock_blob(blob_name,
'hash') # now with hash
client_mock.get_blob_properties = get_blob_properties
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# cleanup
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
self.assertEqual({}, host.policies)
# run
host.update_policies()
# both policies were loaded
self.assertEqual(1, len(host.policies.items()))
# jobs were created
jobs = host.scheduler.get_jobs()
self.assertEqual(1, len([j for j in jobs if j.id == 'blob1.yml']))
def test_update_policies_list_blobs_azure_http_error(
self, get_blob_client_mock, _2, _3, _4, _5):
client_mock = Mock()
client_mock.list_blobs.side_effect = AzureHttpError(
"failed to list blobs", 400)
get_blob_client_mock.return_value = (client_mock, None, None)
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
with self.assertRaises(AzureHttpError):
host.update_policies()
client_mock.list_blobs.assert_called()
def test_poll_queue(self, run_policy_mock, storage_mock, _1, _2, _3):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
host.policies = {
'one': {
'policy':
ContainerHostTest.get_mock_policy({
'name': 'one',
'mode': {
'type': 'container-event',
'events': ['ResourceGroupWrite', 'VnetWrite']
}
})
}
}
q1 = QueueMessage()
q1.id = 1
q1.dequeue_count = 0
q1.content = """eyAgCiAgICJzdWJqZWN0IjoiL3N1YnNjcmlwdGlvbnMvZWE5ODk3NGItNWQyYS00ZDk4LWE3
OGEtMzgyZjM3MTVkMDdlL3Jlc291cmNlR3JvdXBzL3Rlc3RfY29udGFpbmVyX21vZGUiLAogICAiZXZlbnRUeXBlIj
oiTWljcm9zb2Z0LlJlc291cmNlcy5SZXNvdXJjZVdyaXRlU3VjY2VzcyIsCiAgICJldmVudFRpbWUiOiIyMDE5LTA3
LTE2VDE4OjMwOjQzLjM1OTUyNTVaIiwKICAgImlkIjoiNjE5ZDI2NzQtYjM5Ni00MzU2LTk2MTktNmM1YTUyZmU0ZT
g4IiwKICAgImRhdGEiOnsgICAgICAgIAogICAgICAiY29ycmVsYXRpb25JZCI6IjdkZDVhNDc2LWUwNTItNDBlMi05
OWU0LWJiOTg1MmRjMWY4NiIsCiAgICAgICJyZXNvdXJjZVByb3ZpZGVyIjoiTWljcm9zb2Z0LlJlc291cmNlcyIsCi
AgICAgICJyZXNvdXJjZVVyaSI6Ii9zdWJzY3JpcHRpb25zL2VhOTg5NzRiLTVkMmEtNGQ5OC1hNzhhLTM4MmYzNzE1
ZDA3ZS9yZXNvdXJjZUdyb3Vwcy90ZXN0X2NvbnRhaW5lcl9tb2RlIiwKICAgICAgIm9wZXJhdGlvbk5hbWUiOiJNaW
Nyb3NvZnQuUmVzb3VyY2VzL3N1YnNjcmlwdGlvbnMvcmVzb3VyY2VHcm91cHMvd3JpdGUiLAogICAgICAic3RhdHVz
IjoiU3VjY2VlZGVkIiwKICAgfSwKICAgInRvcGljIjoiL3N1YnNjcmlwdGlvbnMvYWE5ODk3NGItNWQyYS00ZDk4LW
E3OGEtMzgyZjM3MTVkMDdlIgp9"""
q2 = QueueMessage()
q2.id = 2
q2.dequeue_count = 0
q2.content = q1.content
# Return 2 messages on first call, then none
storage_mock.get_queue_messages.side_effect = [[q1, q2], []]
host.poll_queue()
self.assertEqual(2, run_policy_mock.call_count)
run_policy_mock.reset_mock()
# Return 5 messages on first call, then 2, then 0
storage_mock.get_queue_messages.side_effect = [[q1, q1, q1, q1, q1],
[q1, q2], []]
host.poll_queue()
self.assertEqual(7, run_policy_mock.call_count)
run_policy_mock.reset_mock()
# High dequeue count
q1.dequeue_count = 100
storage_mock.get_queue_messages.side_effect = [[q1, q2], []]
host.poll_queue()
self.assertEqual(1, run_policy_mock.call_count)
def test_update_policies_missing_mode(self, get_blob_client_mock, _2, _3,
_4, _5):
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob('no-mode-blob.yaml',
'no-mode-blob')
]
client_mock.get_blob_to_path = self.download_missing_mode_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
host.update_policies()
self.assertEqual(1, len(host.policies.items()))
def test_has_required_params(self):
with patch.dict(os.environ, {
constants.ENV_CONTAINER_POLICY_STORAGE: 'foo',
constants.ENV_CONTAINER_EVENT_QUEUE_NAME: 'foo',
constants.ENV_CONTAINER_EVENT_QUEUE_ID: 'foo'
},
clear=False):
self.assertTrue(Host.has_required_params())
with patch.dict(os.environ, {
constants.ENV_CONTAINER_POLICY_STORAGE: 'foo',
constants.ENV_CONTAINER_EVENT_QUEUE_ID: 'foo'
},
clear=False):
self.assertFalse(Host.has_required_params())
def test_unload_policy_file_with_yaml_error(self, os_unlink,
yaml_safe_load, _1, _2, _3, _4,
_5):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# Create a bad yaml file
file_path = tempfile.mktemp(suffix=".yaml")
with open(file_path, 'w') as f:
f.write("bad yaml file")
host.unload_policy_file(file_path, None)
os_unlink.assert_called()
# Clean up the file
os.remove(file_path)
def test_build_options(self):
result = Host.build_options(output_dir='/test/dir',
log_group='test_log_group',
metrics='test_metrics')
self.assertEqual('test_log_group', result['log_group'])
self.assertEqual('/test/dir', result['output_dir'])
self.assertEqual('test_metrics', result['metrics'])
def test_init(self, _1, _2, _3, _4, _5):
host = Host()
jobs = host.scheduler.get_jobs()
update_policy_job = [j for j in jobs if j.id == 'update_policies']
poll_queue_job = [j for j in jobs if j.id == 'poll_queue']
self.assertEqual('test_path', host.policy_cache)
self.assertEqual(2, len(jobs))
self.assertIsNotNone(update_policy_job)
self.assertIsNotNone(poll_queue_job)
def test_init(self, _1, _2, _3, _4):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME, DEFAULT_POLICY_STORAGE)
jobs = host.scheduler.get_jobs()
update_policy_job = [j for j in jobs if j.id == 'update_policies']
poll_queue_job = [j for j in jobs if j.id == 'poll_queue']
self.assertEqual('test_path', host.policy_cache)
self.assertEqual(2, len(jobs))
self.assertIsNotNone(update_policy_job)
self.assertIsNotNone(poll_queue_job)
def test_update_policies_in_subfolders(self, makedirs_mock,
get_blob_client_mock, _2, _3, _4,
_5):
blob_name = "path/to/blob.yml"
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob(blob_name, 'blob')
]
client_mock.get_blob_to_path = self.download_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
host.update_policies()
self.assertEqual(1, len(host.policies.items()))
jobs = host.scheduler.get_jobs()
self.assertEqual(1, len([j for j in jobs if j.id == blob_name]))
blob_dir = os.path.join(host.policy_cache, blob_name)
makedirs_mock.assert_any_call(os.path.dirname(blob_dir))
def test_run_policy_for_event(self, add_job_mock, _0, _1, _2, _3):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
host.policies = {
'one': {
'policy':
ContainerHostTest.get_mock_policy({
'name': 'one',
'mode': {
'type': 'container-event',
'events': ['ResourceGroupWrite', 'VnetWrite']
}
})
}
}
message = QueueMessage()
message.id = 1
message.dequeue_count = 0
message.content = \
"""eyAgCiAgICJzdWJqZWN0IjoiL3N1YnNjcmlwdGlvbnMvZWE5ODk3NGItNWQyYS00ZDk4LWE3OGEt
MzgyZjM3MTVkMDdlL3Jlc291cmNlR3JvdXBzL3Rlc3RfY29udGFpbmVyX21vZGUiLAogICAiZXZl
bnRUeXBlIjoiTWljcm9zb2Z0LlJlc291cmNlcy5SZXNvdXJjZVdyaXRlU3VjY2VzcyIsCiAgICJl
dmVudFRpbWUiOiIyMDE5LTA3LTE2VDE4OjMwOjQzLjM1OTUyNTVaIiwKICAgImlkIjoiNjE5ZDI2
NzQtYjM5Ni00MzU2LTk2MTktNmM1YTUyZmU0ZTg4IiwKICAgImRhdGEiOnsgICAgICAgIAogICAg
ICAiY29ycmVsYXRpb25JZCI6IjdkZDVhNDc2LWUwNTItNDBlMi05OWU0LWJiOTg1MmRjMWY4NiIs
CiAgICAgICJyZXNvdXJjZVByb3ZpZGVyIjoiTWljcm9zb2Z0LlJlc291cmNlcyIsCiAgICAgICJy
ZXNvdXJjZVVyaSI6Ii9zdWJzY3JpcHRpb25zL2VhOTg5NzRiLTVkMmEtNGQ5OC1hNzhhLTM4MmYz
NzE1ZDA3ZS9yZXNvdXJjZUdyb3Vwcy90ZXN0X2NvbnRhaW5lcl9tb2RlIiwKICAgICAgIm9wZXJh
dGlvbk5hbWUiOiJNaWNyb3NvZnQuUmVzb3VyY2VzL3N1YnNjcmlwdGlvbnMvcmVzb3VyY2VHcm91
cHMvd3JpdGUiLAogICAgICAic3RhdHVzIjoiU3VjY2VlZGVkIgogICB9LAogICAidG9waWMiOiIv
c3Vic2NyaXB0aW9ucy9hYTk4OTc0Yi01ZDJhLTRkOTgtYTc4YS0zODJmMzcxNWQwN2UiCn0="""
# run with real match
host.run_policies_for_event(message)
add_job_mock.assert_called_with(
ANY,
id='one619d2674-b396-4356-9619-6c5a52fe4e88',
name='one',
args=ANY,
misfire_grace_time=ANY)
add_job_mock.reset_mock()
# run with no match
host.policies = {}
host.run_policies_for_event(message)
self.assertFalse(add_job_mock.called)
def test_build_options(self):
with patch.dict(os.environ, {
constants.ENV_CONTAINER_OPTION_OUTPUT_DIR: '/test/dir',
constants.ENV_CONTAINER_OPTION_LOG_GROUP: 'test_log_group',
constants.ENV_CONTAINER_OPTION_METRICS: 'test_metrics'
},
clear=False):
result = Host.build_options()
self.assertEqual('test_log_group', result['log_group'])
self.assertEqual('/test/dir', result['output_dir'])
self.assertEqual('test_metrics', result['metrics'])
def test_update_event_subscriptions(self, event_filter_mock, _0, _1, _2,
_3, _4):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
host.event_queue_name = 'testq'
host.policies = {
'one': {
'policy':
ContainerHostTest.get_mock_policy({
'name': 'one',
'mode': {
'type': 'container-event',
'events': ['ResourceGroupWrite', 'VnetWrite']
}
})
},
'two': {
'policy':
ContainerHostTest.get_mock_policy({
'name': 'two',
'mode': {
'type': 'container-event',
'events': ['ResourceGroupWrite']
}
})
},
'three': {
'policy':
ContainerHostTest.get_mock_policy({
'name': 'three',
'mode': {
'type':
'container-event',
'events': [{
'resourceProvider': 'Microsoft.KeyVault/vaults',
'event': 'write'
}]
}
})
}
}
# Verify we get all three events with no duplicates
host.update_event_subscriptions()
event_filter_mock.assert_called_with(
key='Data.OperationName',
values={
'Microsoft.KeyVault/vaults/write',
'Microsoft.Network/virtualNetworks/write',
'Microsoft.Resources/subscriptions/resourceGroups/write'
})
def test_unload_policy_file_without_name(self, os_unlink, yaml_safe_load,
_1, _2, _3, _4):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# Create a bad yaml file (no name field)
file_path = tempfile.mktemp(suffix=".yaml")
with open(file_path, 'w') as f:
f.write("""
policies:
- mode:
type: container-periodic
schedule: '* * * * *'
resource: azure.resourcegroup
""")
host.load_policy(file_path, host.policies)
host.unload_policy_file(file_path, host.policies)
os_unlink.assert_called()
# Clean up the file
os.remove(file_path)
def test_run_policy_handles_exceptions(self, _1, _2, _3, _4, _5):
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
mock_policy = Mock()
mock_policy.push.side_effect = Exception()
host.run_policy(mock_policy, None, None)
def test_init_different_storage_subscription(self, _1, _2, _3, _4, _5):
host = Host(DEFAULT_EVENT_QUEUE_STORAGE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
self.assertIsNot(host.storage_session, host.session)
def test_update_policies_add_remove(self, get_blob_client_mock, _1, _2, _3,
_4):
"""
Run a series of add/update/removal of policy blobs
and verify jobs and caches are updated correctly
"""
# mock blob list call
client_mock = Mock()
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", "hash1")
]
client_mock.get_blob_to_path = self.download_policy_blob
get_blob_client_mock.return_value = (client_mock, None, None)
# init
host = Host(DEFAULT_EVENT_QUEUE_ID, DEFAULT_EVENT_QUEUE_NAME,
DEFAULT_POLICY_STORAGE)
# cleanup
self.addCleanup(lambda: shutil.rmtree(host.policy_cache))
self.assertEqual({}, host.policies)
# Initial load
host.update_policies()
self.assertEqual(1, len(host.policies.items()))
##################
# Add two policies
##################
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", "hash1"),
ContainerHostTest.get_mock_blob("blob2.yml", "hash2"),
ContainerHostTest.get_mock_blob("blob3.yml", "hash3")
]
host.update_policies()
self.assertEqual(3, len(host.policies.items()))
self.assertIsNotNone(host.policies['blob1.yml'])
self.assertIsNotNone(host.policies['blob2.yml'])
self.assertIsNotNone(host.policies['blob3.yml'])
# jobs were updated
jobs = host.scheduler.get_jobs()
self.assertEqual(3,
len([j for j in jobs if j.func == host.run_policy]))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob1.yml']))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob2.yml']))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob3.yml']))
##############################################
# Add one, remove one, update one
##############################################
client_mock.list_blobs.return_value = [
ContainerHostTest.get_mock_blob("blob1.yml", "hash1"),
ContainerHostTest.get_mock_blob("blob4.yml", "hash4"),
ContainerHostTest.get_mock_blob("blob3.yml", "hash3_new")
]
host.update_policies()
self.assertEqual(3, len(host.policies.items()))
self.assertIsNotNone(host.policies['blob1.yml'])
self.assertIsNotNone(host.policies['blob4.yml'])
self.assertIsNotNone(host.policies['blob3.yml'])
self.assertEqual('hash3_new', host.blob_cache['blob3.yml'])
# jobs were updated
jobs = host.scheduler.get_jobs()
self.assertEqual(3,
len([j for j in jobs if j.func == host.run_policy]))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob1.yml']))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob4.yml']))
self.assertEqual(1, len([j for j in jobs if j.id == 'blob3.yml']))
############
# remove all
############
client_mock.list_blobs.return_value = []
host.update_policies()
self.assertEqual(0, len(host.policies.items()))
# jobs were updated
jobs = host.scheduler.get_jobs()
self.assertEqual(0,
len([j for j in jobs if j.func == host.run_policy]))
def test_build_options_empty(self, _):
result = Host.build_options()
self.assertEqual(None, result['log_group'])
self.assertEqual('test_path', result['output_dir'])
self.assertEqual(None, result['metrics'])