def test_eq(self):
flow_record = FlowRecord.from_cwl_event({'message': V2_RECORDS[0]})
equal_record = FlowRecord.from_cwl_event({'message': V2_RECORDS[0]})
unequal_record = FlowRecord.from_cwl_event({'message': V2_RECORDS[1]})
self.assertEqual(flow_record, equal_record)
self.assertNotEqual(flow_record, unequal_record)
self.assertNotEqual(flow_record, Ellipsis)
def test_aggregated_records_custom(self):
# Aggregate by interface_id
events = [
{'message': V2_RECORDS[1]},
{'message': V2_RECORDS[2].replace('REJECT', 'ACCEPT')},
]
all_records = (FlowRecord.from_cwl_event(x) for x in events)
key_fields = ('interface_id', 'srcaddr', 'srcport', 'dstport')
results = aggregated_records(all_records, key_fields=key_fields)
actual = sorted(results, key=lambda x: x['interface_id'])
expected = [
{
'srcaddr': '192.0.2.1',
'srcport': 49152,
'interface_id': 'eni-102010ab',
'dstport': 443,
'start': datetime(2015, 8, 12, 13, 47, 44),
'end': datetime(2015, 8, 12, 13, 47, 45),
'packets': 20,
'bytes': 1680,
},
{
'srcaddr': '192.0.2.1',
'srcport': 49152,
'interface_id': 'eni-102010cd',
'dstport': 443,
'start': datetime(2015, 8, 12, 13, 47, 43),
'end': datetime(2015, 8, 12, 13, 47, 46),
'packets': 20,
'bytes': 1680,
},
]
self.assertEqual(actual, expected)
def test_iteration(self):
paginator = MagicMock()
paginator.paginate.return_value = [
{
'events': [
{'logStreamName': 'log_0', 'message': V2_RECORDS[0]},
{'logStreamName': 'log_0', 'message': V2_RECORDS[1]},
],
},
{
'events': [
{'logStreamName': 'log_0', 'message': V2_RECORDS[2]},
{'logStreamName': 'log_1', 'message': V2_RECORDS[3]},
{'logStreamName': 'log_2', 'message': V2_RECORDS[4]},
],
},
]
self.mock_client.get_paginator.return_value = paginator
# Calling list on the instance causes it to iterate through all records
actual = [next(self.inst)] + list(self.inst)
expected = [
FlowRecord.from_cwl_event({'message': x}) for x in V2_RECORDS
]
self.assertEqual(actual, expected)
expected_bytes = 0
all_pages = paginator.paginate.return_value
expected_bytes = sum(
len(e['message']) for p in all_pages for e in p['events']
)
self.assertEqual(self.inst.bytes_processed, expected_bytes)
def test_millisecond_timestamp(self):
# This record has millisecond timestamps
record = ('2 123456789010 eni-4b118871 - - - - - - - '
'1512564058000 1512564059000 - SKIPDATA')
flow_record = FlowRecord.from_cwl_event({'message': record})
self.assertEqual(flow_record.start, datetime(2017, 12, 6, 12, 40, 58))
self.assertEqual(flow_record.end, datetime(2017, 12, 6, 12, 40, 59))
def test_str(self):
flow_record = FlowRecord.from_cwl_event({'message': V2_RECORDS[0]})
actual = str(flow_record)
expected = (
'version: 2, account_id: 123456789010, '
'interface_id: eni-102010ab, srcaddr: 198.51.100.1, '
'dstaddr: 192.0.2.1, srcport: 443, dstport: 49152, protocol: 6, '
'packets: 10, bytes: 840, start: 2015-08-12 13:47:43, '
'end: 2015-08-12 13:47:44, action: ACCEPT, log_status: OK')
self.assertEqual(actual, expected)
def test_hash(self):
record_set = {
FlowRecord.from_cwl_event({'message': V2_RECORDS[0]}),
FlowRecord.from_cwl_event({'message': V2_RECORDS[0]}),
FlowRecord.from_cwl_event({'message': V2_RECORDS[1]}),
FlowRecord.from_cwl_event({'message': V2_RECORDS[1]}),
FlowRecord.from_cwl_event({'message': V2_RECORDS[2]}),
FlowRecord.from_cwl_event({'message': V2_RECORDS[2]}),
}
self.assertEqual(len(record_set), 3)
def test_aggregated_records(self):
# Aggregate by 5-tuple by default
events = [
{
'message': V2_RECORDS[0]
},
{
'message': V2_RECORDS[1]
},
{
'message': V2_RECORDS[2].replace('REJECT', 'ACCEPT')
},
{
'message': V2_RECORDS[3]
},
]
all_records = (FlowRecord.from_cwl_event(x) for x in events)
results = aggregated_records(all_records)
actual = sorted(results, key=lambda x: x['srcaddr'])
expected = [
{
'srcaddr': '192.0.2.1',
'srcport': 49152,
'dstaddr': '198.51.100.1',
'dstport': 443,
'protocol': 6,
'start': datetime(2015, 8, 12, 13, 47, 43),
'end': datetime(2015, 8, 12, 13, 47, 46),
'packets': 40,
'bytes': 3360,
},
{
'srcaddr': '198.51.100.1',
'srcport': 443,
'dstaddr': '192.0.2.1',
'dstport': 49152,
'protocol': 6,
'start': datetime(2015, 8, 12, 13, 47, 43),
'end': datetime(2015, 8, 12, 13, 47, 44),
'packets': 10,
'bytes': 840,
},
]
self.assertEqual(actual, expected)
def test_iteration_error(self):
# Simulate the paginator failing
def _get_paginator(*args, **kwargs):
event_0 = {'logStreamName': 'log_0', 'message': V2_RECORDS[0]}
event_1 = {'logStreamName': 'log_0', 'message': V2_RECORDS[1]}
for item in [{'events': [event_0, event_1]}]:
yield item
err_msg = '{}: {}'.format(DUPLICATE_NEXT_TOKEN_MESSAGE, 'token')
raise PaginationError(message=err_msg)
self.mock_client.get_paginator.return_value.paginate.side_effect = (
_get_paginator)
# Don't fail if botocore's paginator raises a PaginationError
actual = [next(self.inst)] + list(self.inst)
records = V2_RECORDS[:2]
expected = [FlowRecord.from_cwl_event({'message': x}) for x in records]
self.assertEqual(actual, expected)
def test_to_dict(self):
flow_record = FlowRecord.from_cwl_event({'message': V2_RECORDS[2]})
actual = flow_record.to_dict()
expected = {
'account_id': '123456789010',
'action': 'REJECT',
'bytes': 1680,
'dstaddr': '198.51.100.1',
'dstport': 443,
'end': datetime(2015, 8, 12, 13, 47, 46),
'interface_id': 'eni-102010cd',
'log_status': 'OK',
'packets': 20,
'protocol': 6,
'srcaddr': '192.0.2.1',
'srcport': 49152,
'start': datetime(2015, 8, 12, 13, 47, 43),
'version': 2,
}
self.assertEqual(actual, expected)
'49152 443 6 20 1680 1439387265 1439387266 REJECT OK '
'- - - - - - -'
),
(
'2 123456789010 eni-1a2b3c4d - - - - - - - '
'1431280876 1431280934 - NODATA '
'- - - - - - -'
),
(
'2 123456789010 eni-4b118871 - - - - - - - '
'1431280876 1431280934 - SKIPDATA '
'- - - - - - -'
),
]
SAMPLE_RECORDS = [
FlowRecord.from_cwl_event({'message': m}) for m in SAMPLE_INPUT
]
class MainTestCase(TestCase):
@patch('flowlogs_reader.__main__.FlowLogsReader', autospec=True)
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'])
