def test_main_aggregate(self, mock_reader):
mock_reader.return_value = [SAMPLE_RECORDS[0], SAMPLE_RECORDS[0]]
with patch('sys.stdout', new_callable=StringIO) as mock_stdout:
main(['mygroup', 'aggregate'])
output = mock_stdout.getvalue().splitlines()
actual_header = output[0].split('\t')
expected_header = [
'bytes',
'dstaddr',
'dstport',
'end',
'packets',
'protocol',
'srcaddr',
'srcport',
'start',
]
self.assertEqual(actual_header, expected_header)
actual_line = output[1].split('\t')
expected_line = [
'1680',
'192.0.2.1',
'49152',
'2015-08-12 13:47:44',
'20',
'6',
'198.51.100.1',
'443',
'2015-08-12 13:47:43',
]
self.assertEqual(actual_line, expected_line)
def test_main_assume_role_s3(self, mock_boto3, mock_reader):
mock_boto3.client.return_value.assume_role.return_value = {
'Credentials': {
'AccessKeyId': 'myaccesskeyid',
'SecretAccessKey': 'mysecretaccesskey',
'SessionToken': 'mysessiontoken',
}
}
mock_client = MagicMock()
mock_boto3.session.Session.return_value.client.return_value = (
mock_client
)
mock_reader.return_value = []
args = [
'--role-arn', 'myarn',
'--external-id', 'uuid4',
'--location-type', 's3',
'mybucket'
]
main(args)
session = mock_boto3.session.Session
session.assert_called_once_with(
aws_access_key_id='myaccesskeyid',
aws_secret_access_key='mysecretaccesskey',
aws_session_token='mysessiontoken',
)
session.return_value.client.assert_called_once_with('s3')
mock_reader.assert_called_once_with(
'mybucket', boto_client=mock_client
)
def test_main(self, mock_reader):
main(['mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-east-1'
)
main(['-s', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-east-1',
start_time=datetime(2015, 5, 5, 14, 20)
)
main(['--end-time', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-east-1',
end_time=datetime(2015, 5, 5, 14, 20)
)
main([
'--time-format', '%Y-%m-%d',
'--start-time', '2015-05-05',
'mygroup'
])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-east-1',
start_time=datetime(2015, 5, 5)
)
main(['--region', 'us-west-1', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-west-1'
)
def test_main_assume_role(self, mock_boto3, mock_reader):
mock_boto3.client.return_value.assume_role.return_value = {
'Credentials': {
'AccessKeyId': 'myaccesskeyid',
'SecretAccessKey': 'mysecretaccesskey',
'SessionToken': 'mysessiontoken',
}
}
mock_client = MagicMock()
mock_boto3.session.Session.return_value.client.return_value = (
mock_client
)
mock_reader.return_value = []
main(['--role-arn', 'myarn', '--external-id', 'uuid4', 'mygroup'])
session = mock_boto3.session.Session
session.assert_called_once_with(
aws_access_key_id='myaccesskeyid',
aws_secret_access_key='mysecretaccesskey',
aws_session_token='mysessiontoken',
)
session.return_value.client.assert_called_once_with('logs')
mock_reader.assert_called_once_with(
log_group_name='mygroup', boto_client=mock_client
)
def test_main_aggregate(self, mock_reader):
mock_reader.return_value = [SAMPLE_RECORDS[0], SAMPLE_RECORDS[0]]
with patch('sys.stdout', new_callable=StringIO) as mock_stdout:
main(['mygroup', 'aggregate'])
output = mock_stdout.getvalue().splitlines()
actual_header = output[0].split('\t')
expected_header = [
'bytes',
'dstaddr',
'dstport',
'end',
'packets',
'protocol',
'srcaddr',
'srcport',
'start',
]
self.assertEqual(actual_header, expected_header)
actual_line = output[1].split('\t')
expected_line = [
'1680',
'192.0.2.1',
'49152',
'2015-08-12 13:47:44',
'20',
'6',
'198.51.100.1',
'443',
'2015-08-12 13:47:43',
]
self.assertEqual(actual_line, expected_line)
def test_main_print(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup'])
for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_print(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup'])
for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_findip(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', 'findip', '198.51.100.2'])
expected_result = [SAMPLE_INPUT[2]]
for call, record in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_findip(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', 'findip', '198.51.100.2'])
expected_result = [SAMPLE_INPUT[2]]
for call, record in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_missing_arn(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['--external-id', 'uuid4', 'mygroup'])
expected_result = [
'must give a --role-arn if an --external-id is given',
]
for call, result in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, result)
def test_main_missing_arn(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['--external-id', 'uuid4', 'mygroup'])
expected_result = [
'must give a --role-arn if an --external-id is given',
]
for call, result in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, result)
def test_main_bad_action(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', '__'])
expected_result = [
'unknown action: __',
'known actions: {}'.format(', '.join(actions)),
]
for call, result in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, result)
def test_main_bad_action(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', '__'])
expected_result = [
'unknown action: __',
'known actions: {}'.format(', '.join(actions)),
]
for call, result in zip_longest(mock_out.mock_calls, expected_result):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, result)
def test_s3_destination(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(
[
'mybucket/myprefix',
'--location-type', 's3',
'--include-accounts', '999999999998, 999999999999',
'--include-regions', 'us-east-1,us-east-2',
]
)
mock_reader.assert_called_once_with(
location='mybucket/myprefix',
include_accounts=['999999999998', '999999999999'],
include_regions=['us-east-1', 'us-east-2'],
)
for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_ipset(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', 'ipset'])
expected_set = set()
for record in SAMPLE_INPUT:
data = record.split()
expected_set.add(data[3])
expected_set.add(data[4])
# don't include SKIPDATA/NODATA in results
expected_set.remove('-')
# make sure the number of lines are the same as the size of the set
self.assertEqual(len(mock_out.mock_calls), len(expected_set))
actual_set = set()
for __, args, kwargs in mock_out.mock_calls:
line = args[0]
actual_set.add(line)
self.assertEqual(actual_set, expected_set)
def test_main_ipset(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
main(['mygroup', 'ipset'])
expected_set = set()
for record in SAMPLE_INPUT:
data = record.split()
expected_set.add(data[3])
expected_set.add(data[4])
# don't include SKIPDATA/NODATA in results
expected_set.remove('-')
# make sure the number of lines are the same as the size of the set
self.assertEqual(len(mock_out.mock_calls), len(expected_set))
actual_set = set()
for __, args, kwargs in mock_out.mock_calls:
line = args[0]
actual_set.add(line)
self.assertEqual(actual_set, expected_set)
def test_main_assume_role(self, mock_boto3, mock_reader):
mock_boto3.client.return_value.assume_role.return_value = {
'Credentials': {
'AccessKeyId': 'myaccesskeyid',
'SecretAccessKey': 'mysecretaccesskey',
'SessionToken': 'mysessiontoken',
}
}
mock_client = MagicMock()
mock_boto3.session.Session.return_value.client.return_value = (
mock_client
)
mock_reader.return_value = []
main(['--role-arn', 'myarn', '--external-id', 'uuid4', 'mygroup'])
mock_boto3.session.Session.assert_called_once_with(
aws_access_key_id='myaccesskeyid',
aws_secret_access_key='mysecretaccesskey',
aws_session_token='mysessiontoken',
)
mock_reader.assert_called_once_with(
log_group_name='mygroup', boto_client=mock_client
)
def test_main_print_count(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
with self.assertRaises(ValueError):
main(['mygroup', 'print', 'two'])
with self.assertRaises(RuntimeError):
main(['mygroup', 'print', '2', '3'])
main(['mygroup', 'print', '2'])
for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT[:2]):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main_print_count(self, mock_out, mock_reader):
mock_out.stdout = io.BytesIO()
mock_reader.return_value = SAMPLE_RECORDS
with self.assertRaises(ValueError):
main(['mygroup', 'print', 'two'])
with self.assertRaises(RuntimeError):
main(['mygroup', 'print', '2', '3'])
main(['mygroup', 'print', '2'])
for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT[:2]):
__, args, kwargs = call
line = args[0]
self.assertEqual(line, record)
def test_main(self, mock_reader):
main(['mygroup'])
mock_reader.assert_called_with(log_group_name='mygroup')
main(['-s', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', start_time=datetime(2015, 5, 5, 14, 20),
)
main(['--end-time', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', end_time=datetime(2015, 5, 5, 14, 20),
)
main([
'--time-format', '%Y-%m-%d',
'--start-time', '2015-05-05',
'mygroup'
])
mock_reader.assert_called_with(
log_group_name='mygroup', start_time=datetime(2015, 5, 5),
)
main(['--region', 'us-west-1', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-west-1',
)
main(['--profile', 'my-profile', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', profile_name='my-profile'
)
main(['--filter-pattern', 'REJECT', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', filter_pattern='REJECT'
)
def test_main(self, mock_reader):
main(['mygroup'])
mock_reader.assert_called_with(log_group_name='mygroup', fields=None)
main(['-s', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup',
fields=None,
start_time=datetime(2015, 5, 5, 14, 20),
)
main(['--end-time', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup',
fields=None,
end_time=datetime(2015, 5, 5, 14, 20),
)
main([
'--time-format',
'%Y-%m-%d',
'--start-time',
'2015-05-05',
'mygroup',
'--thread-count',
'2',
])
mock_reader.assert_called_with(
log_group_name='mygroup',
fields=None,
start_time=datetime(2015, 5, 5),
thread_count=2,
)
main(['--region', 'us-west-1', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup',
fields=None,
region_name='us-west-1',
)
main(['--profile', 'my-profile', 'mygroup'])
mock_reader.assert_called_with(log_group_name='mygroup',
fields=None,
profile_name='my-profile')
main([
'--filter-pattern',
'REJECT',
'--fields',
'${account-id} ${action}',
'mygroup',
])
mock_reader.assert_called_with(
log_group_name='mygroup',
fields=('account-id', 'action'),
filter_pattern='REJECT',
)
def test_main(self, mock_reader):
main(['mygroup'])
mock_reader.assert_called_with(log_group_name='mygroup')
main(['-s', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', start_time=datetime(2015, 5, 5, 14, 20),
)
main(['--end-time', '2015-05-05 14:20:00', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', end_time=datetime(2015, 5, 5, 14, 20),
)
main([
'--time-format', '%Y-%m-%d',
'--start-time', '2015-05-05',
'mygroup'
])
mock_reader.assert_called_with(
log_group_name='mygroup', start_time=datetime(2015, 5, 5),
)
main(['--region', 'us-west-1', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', region_name='us-west-1',
)
main(['--profile', 'my-profile', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', profile_name='my-profile'
)
main(['--filter-pattern', 'REJECT', 'mygroup'])
mock_reader.assert_called_with(
log_group_name='mygroup', filter_pattern='REJECT'
)
