예제 #1
0
    def test_region_name(self, mock_client):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_client.assert_called_with('logs', region_name='some-region')

        # None specified
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs')
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()
예제 #3
0
    def setUp(self, mock_boto3):
        self.mock_client = MagicMock()
        mock_boto3.client.return_value = self.mock_client

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
        )
예제 #4
0
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )
예제 #5
0
    def test_region(self, mock_client):
        # Region specified
        FlowLogsReader('some_group', region_name='some-region')
        mock_client.assert_called_with('logs', region_name='some-region')

        # No region specified - assume configuration file worked
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs')

        # No region specified and no configuration file
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_client.side_effect = mock_response
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs', region_name=DEFAULT_REGION_NAME)
    def setUp(self, mock_boto3):
        self.mock_client = MagicMock()
        mock_boto3.client.return_value = self.mock_client

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
        )
예제 #7
0
    def test_region_name(self, mock_session):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # Region specified for client, not for session
        FlowLogsReader('some_group',
                       boto_client_kwargs={'region_name': 'my-region'})
        mock_session.Session().client.assert_called_with(
            'logs', region_name='my-region')

        # No region specified for session or client - use the default
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_session.Session().client.side_effect = mock_response

        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME)
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )
예제 #9
0
 def test_get_fields(self, mock_client):
     cwl_client = MagicMock()
     ec2_client = mock_client.return_value
     ec2_client.describe_flow_logs.return_value = {
         'FlowLogs': [
             {'LogFormat': '${srcaddr} ${dstaddr} ${start} ${log-status}'}
         ]
     }
     reader = FlowLogsReader(
         'some_group',
         boto_client=cwl_client,
         fields=None,
     )
     self.assertEqual(
         reader.fields, ('srcaddr', 'dstaddr', 'start', 'log_status')
     )
     ec2_client.describe_flow_logs.assert_called_once_with(
         Filters=[{'Name': 'log-group-name', 'Values': ['some_group']}]
     )
예제 #10
0
class FlowLogsReaderTestCase(TestCase):
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(datetime.utcfromtimestamp(self.inst.start_ms // 1000),
                         self.start_time)

        self.assertEqual(datetime.utcfromtimestamp(self.inst.end_ms // 1000),
                         self.end_time)

        self.assertEqual(self.inst.paginator_kwargs['filterPattern'], 'REJECT')

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_region_name(self, mock_session):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # Region specified for client, not for session
        FlowLogsReader('some_group',
                       boto_client_kwargs={'region_name': 'my-region'})
        mock_session.Session().client.assert_called_with(
            'logs', region_name='my-region')

        # No region specified for session or client - use the default
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_session.Session().client.side_effect = mock_response

        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME)

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

    def test_read_streams(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [0]
            },
            {
                'events': [1, 2]
            },
            {
                'events': [3, 4, 5]
            },
        ]

        self.mock_client.get_paginator.return_value = paginator

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        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)

    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_iteration_unexpecetd_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': V2_RECORDS[0]}
            yield {'events': [event_0]}
            raise PaginationError(message='other error')

        self.mock_client.get_paginator.return_value.paginate.side_effect = (
            _get_paginator)

        # Fail for unexpected PaginationError
        self.assertRaises(PaginationError, lambda: list(self.inst))

    def test_threads(self):
        inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
            thread_count=1,
        )

        paginators = []

        def _get_paginator(operation):
            nonlocal paginators

            paginator = MagicMock()
            if operation == 'describe_log_streams':
                paginator.paginate.return_value = [
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'too_late',
                                'firstEventTimestamp': inst.end_ms,
                                'lastEventTimestamp': inst.start_ms,
                            },
                            {
                                'logStreamName':
                                'too_late',
                                'firstEventTimestamp':
                                inst.end_ms - 1,
                                'lastEventTimestamp':
                                (inst.start_ms - LAST_EVENT_DELAY_MSEC - 1),
                            },
                        ],
                    },
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'first_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                            {
                                'logStreamName': 'second_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                        ],
                    },
                ]
            elif operation == 'filter_log_events':
                paginator.paginate.return_value = [
                    {
                        'events': [
                            {
                                'message': V2_RECORDS[0]
                            },
                            {
                                'message': V2_RECORDS[1]
                            },
                        ],
                    },
                    {
                        'events': [
                            {
                                'message': V2_RECORDS[2]
                            },
                            {
                                'message': V2_RECORDS[3]
                            },
                        ],
                    },
                ]
            else:
                self.fail('invalid operation')

            paginators.append(paginator)
            return paginator

        self.mock_client.get_paginator.side_effect = _get_paginator
        events = list(inst)
        self.assertEqual(len(events), 8)

        paginators[0].paginate.assert_called_once_with(
            logGroupName='group_name',
            orderBy='LastEventTime',
            descending=True,
        )
        paginators[1].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['first_stream'],
        )
        paginators[2].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['second_stream'],
        )
class FlowLogsReaderTestCase(TestCase):
    @patch('flowlogs_reader.flowlogs_reader.boto3', autospec=True)
    def setUp(self, mock_boto3):
        self.mock_session = MagicMock()
        mock_boto3.session.Session.return_value = self.mock_session
        self.mock_client = MagicMock()
        self.mock_session.client.return_value = self.mock_client

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(datetime.utcfromtimestamp(self.inst.start_ms // 1000),
                         self.start_time)

        self.assertEqual(datetime.utcfromtimestamp(self.inst.end_ms // 1000),
                         self.end_time)

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_region_name(self, mock_session):
        # Region specified
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # No region specified - assume configuration file worked
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

        # Region specified in boto_client_kwargs
        FlowLogsReader('some_group',
                       boto_client_kwargs={'region_name': 'my-region'})
        mock_session.Session.assert_called_with(region_name='my-region')

        # No region specified and no configuration file
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_session.Session().client.side_effect = mock_response
        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME)

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

    def test_read_streams(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [0]
            },
            {
                'events': [1, 2]
            },
            {
                'events': [3, 4, 5]
            },
        ]

        self.mock_client.get_paginator.return_value = paginator

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        self.assertEqual(actual, expected)

    def test_iteration(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[0]
                    },
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[1]
                    },
                ],
            },
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[2]
                    },
                    {
                        'logStreamName': 'log_1',
                        'message': SAMPLE_RECORDS[3]
                    },
                    {
                        'logStreamName': 'log_2',
                        'message': SAMPLE_RECORDS[4]
                    },
                ],
            },
        ]

        self.mock_client.get_paginator.return_value = paginator

        # Calling list on the instance causes it to iterate through all records
        actual = list(self.inst)
        expected = [FlowRecord.from_message(x) for x in SAMPLE_RECORDS]
        self.assertEqual(actual, expected)
class FlowLogsReaderTestCase(TestCase):
    @patch('flowlogs_reader.flowlogs_reader.boto3', autospec=True)
    def setUp(self, mock_boto3):
        self.mock_session = MagicMock()
        mock_boto3.session.Session.return_value = self.mock_session
        self.mock_client = MagicMock()
        self.mock_session.client.return_value = self.mock_client

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.start_ms // 1000),
            self.start_time
        )

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.end_ms // 1000),
            self.end_time
        )

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_region_name(self, mock_session):
        # Region specified
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # No region specified - assume configuration file worked
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

        # Region specified in boto_client_kwargs
        FlowLogsReader(
            'some_group', boto_client_kwargs={'region_name': 'my-region'}
        )
        mock_session.Session.assert_called_with(region_name='my-region')

        # No region specified and no configuration file
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError
        mock_session.Session().client.side_effect = mock_response
        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME
        )

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

    def test_read_streams(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {'events': [0]}, {'events': [1, 2]}, {'events': [3, 4, 5]},
        ]

        self.mock_client.get_paginator.return_value = paginator

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        self.assertEqual(actual, expected)

    def test_iteration(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]},
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[1]},
                ],
            },
            {
                'events': [
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[2]},
                    {'logStreamName': 'log_1', 'message': SAMPLE_RECORDS[3]},
                    {'logStreamName': 'log_2', 'message': SAMPLE_RECORDS[4]},
                ],
            },
        ]

        self.mock_client.get_paginator.return_value = paginator

        # Calling list on the instance causes it to iterate through all records
        actual = list(self.inst)
        expected = [FlowRecord.from_message(x) for x in SAMPLE_RECORDS]
        self.assertEqual(actual, expected)
class FlowLogsReaderTestCase(TestCase):
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.start_ms // 1000),
            self.start_time
        )

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.end_ms // 1000),
            self.end_time
        )

        self.assertEqual(
            self.inst.paginator_kwargs['filterPattern'],
            'REJECT'
        )

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_region_name(self, mock_session):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # Region specified for client, not for session
        FlowLogsReader(
            'some_group', boto_client_kwargs={'region_name': 'my-region'}
        )
        mock_session.Session().client.assert_called_with(
            'logs', region_name='my-region'
        )

        # No region specified for session or client - use the default
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError
        mock_session.Session().client.side_effect = mock_response

        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME
        )

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

    def test_read_streams(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {'events': [0]}, {'events': [1, 2]}, {'events': [3, 4, 5]},
        ]

        self.mock_client.get_paginator.return_value = paginator

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        self.assertEqual(actual, expected)

    def test_iteration(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]},
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[1]},
                ],
            },
            {
                'events': [
                    {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[2]},
                    {'logStreamName': 'log_1', 'message': SAMPLE_RECORDS[3]},
                    {'logStreamName': 'log_2', 'message': SAMPLE_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_message(x) for x in SAMPLE_RECORDS]
        self.assertEqual(actual, expected)

    def test_iteration_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]}
            event_1 = {'logStreamName': 'log_0', 'message': SAMPLE_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)
        expected = [FlowRecord.from_message(x) for x in SAMPLE_RECORDS[:2]]
        self.assertEqual(actual, expected)

    def test_iteration_unexpecetd_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]}
            yield {'events': [event_0]}
            raise PaginationError(message='other error')

        self.mock_client.get_paginator.return_value.paginate.side_effect = (
            _get_paginator
        )

        # Fail for unexpected PaginationError
        self.assertRaises(PaginationError, lambda: list(self.inst))
예제 #14
0
def get_connections(account_id, region, connections=None):
    logging.info('Getting connections for {} {}..'.format(account_id, region))
    sts = boto3.client('sts')
    response = sts.assume_role(RoleArn='arn:aws:iam::' + account_id +
                               ':role/fullstop',
                               RoleSessionName='fullstop',
                               DurationSeconds=900)

    session = boto3.Session(
        aws_access_key_id=response['Credentials']['AccessKeyId'],
        aws_secret_access_key=response['Credentials']['SecretAccessKey'],
        aws_session_token=response['Credentials']['SessionToken'],
        region_name=region)

    ec2_client = session.client('ec2')
    rds = session.client('rds')

    instance_ids = []
    interfaces = {}

    logging.info('%s: Collecting network interfaces..', account_id)
    res = ec2_client.describe_network_interfaces()
    lb_names = []
    for iface in res['NetworkInterfaces']:
        if 'Association' in iface:
            # public IP involved
            interfaces[iface['NetworkInterfaceId']] = iface
            descr = iface.get('Description')
            if descr.startswith('ELB'):
                words = descr.split()
                lb_names.append(words[-1])
        if 'Attachment' in iface and 'InstanceId' in iface['Attachment']:
            instance_ids.append(iface['Attachment']['InstanceId'])

    lb_dns_names = get_lb_dns_names(session, lb_names)

    res = rds.describe_db_instances()
    for db in res['DBInstances']:
        if db['PubliclyAccessible']:
            host, port = (db['Endpoint']['Address'], db['Endpoint']['Port'])
            try:
                ai = socket.getaddrinfo(host,
                                        port,
                                        family=socket.AF_INET,
                                        socktype=socket.SOCK_STREAM)
            except:
                ai = []
            for _, _, _, _, ip_port in ai:
                ip, _ = ip_port
                NAMES[ip] = '/'.join((account_id, region, host))

    local_names = {}
    instance_count = 0
    logging.info('%s: Collecting public EC2 instances..', account_id)
    res = ec2_client.describe_instances(InstanceIds=instance_ids)
    for reservation in res['Reservations']:
        for inst in reservation['Instances']:
            instance_count += 1
            if 'PrivateIpAddress' in inst and 'Tags' in inst:
                name = ''.join(
                    [x['Value'] for x in inst['Tags'] if x['Key'] == 'Name'])
                local_names[inst['PrivateIpAddress']] = '/'.join(
                    (name, inst.get('PublicIpAddress', '')))

    logging.info(
        '%s: Got {} interfaces, {} load balancers and {} instances'.format(
            len(interfaces), len(lb_dns_names), instance_count), account_id)

    connections = collections.Counter() if connections is None else connections
    now = datetime.datetime.utcnow()
    start_time = LAST_TIMES.get(account_id,
                                now - datetime.timedelta(minutes=10))
    reader = FlowLogsReader('vpc-flowgroup',
                            region_name=region,
                            start_time=start_time,
                            end_time=now)
    reader.logs_client = session.client('logs')
    record_count = 0
    new_connections = 0
    for record in reader:
        # just consider accepted packets
        if record.action == 'ACCEPT':
            record_count += 1
            src = ipaddress.ip_address(record.srcaddr)
            # only look at packets received at public interfaces
            if record.interface_id in interfaces and not src.is_private:
                name = get_name(record.srcaddr)
                dest = interfaces.get(record.interface_id,
                                      {}).get('Description')
                if not dest or dest.startswith('Primary'):
                    # EC2 instance
                    if record.srcaddr in AWS_IPS and AWS_S3_DOMAIN_PATTERN.match(
                            name.split('/')[0]):
                        # ignore S3 public IPs
                        # (most probably packets from S3 to public EC2)
                        continue
                    dest = local_names.get(record.dstaddr, record.dstaddr)
                elif dest.startswith('ELB'):
                    # ELB
                    words = dest.split()
                    dest = lb_dns_names.get(words[-1], dest)
                elif dest.startswith('RDS'):
                    # RDS instance
                    public_ip = interfaces.get(record.interface_id,
                                               {}).get('Association',
                                                       {}).get('PublicIp', '')
                    dest = NAMES.get(public_ip, 'RDS/' + public_ip)
                elif dest:
                    dest += '/' + interfaces.get(record.interface_id, {}).get(
                        'Association', {}).get('PublicIp', '')
                if 'NAT' not in dest and 'Odd' not in dest:
                    conn = (name, dest, record.dstport)
                    if conn not in connections:
                        new_connections += 1
                    connections[conn] += 1
    logging.info(
        '%s: Got {} records and {} new connections'.format(
            record_count, new_connections), account_id)
    LAST_TIMES[account_id] = now
    return connections
예제 #15
0
def get_connections(account_id, region, connections=None):
    logging.info('Getting connections for {} {}..'.format(account_id, region))
    sts = boto3.client('sts')
    response = sts.assume_role(RoleArn='arn:aws:iam::' + account_id + ':role/fullstop',
                               RoleSessionName='fullstop',
                               DurationSeconds=900)

    session = boto3.Session(aws_access_key_id=response['Credentials']['AccessKeyId'],
                            aws_secret_access_key=response['Credentials']['SecretAccessKey'],
                            aws_session_token=response['Credentials']['SessionToken'],
                            region_name=region)

    ec2_client = session.client('ec2')
    rds = session.client('rds')

    instance_ids = []
    interfaces = {}

    logging.info('%s: Collecting network interfaces..', account_id)
    res = ec2_client.describe_network_interfaces()
    lb_names = []
    for iface in res['NetworkInterfaces']:
        if 'Association' in iface:
            # public IP involved
            interfaces[iface['NetworkInterfaceId']] = iface
            descr = iface.get('Description')
            if descr.startswith('ELB'):
                words = descr.split()
                lb_names.append(words[-1])
        if 'Attachment' in iface and 'InstanceId' in iface['Attachment']:
            instance_ids.append(iface['Attachment']['InstanceId'])

    lb_dns_names = get_lb_dns_names(session, lb_names)

    res = rds.describe_db_instances()
    for db in res['DBInstances']:
        if db['PubliclyAccessible']:
            host, port = (db['Endpoint']['Address'], db['Endpoint']['Port'])
            try:
                ai = socket.getaddrinfo(host, port, family=socket.AF_INET, socktype=socket.SOCK_STREAM)
            except:
                ai = []
            for _, _, _, _, ip_port in ai:
                ip, _ = ip_port
                NAMES[ip] = '/'.join((account_id, region, host))

    local_names = {}
    instance_count = 0
    logging.info('%s: Collecting public EC2 instances..', account_id)
    res = ec2_client.describe_instances(InstanceIds=instance_ids)
    for reservation in res['Reservations']:
        for inst in reservation['Instances']:
            instance_count += 1
            if 'PrivateIpAddress' in inst and 'Tags' in inst:
                name = ''.join([x['Value'] for x in inst['Tags'] if x['Key'] == 'Name'])
                local_names[inst['PrivateIpAddress']] = '/'.join((name, inst.get('PublicIpAddress', '')))

    logging.info('%s: Got {} interfaces, {} load balancers and {} instances'.format(
                 len(interfaces), len(lb_dns_names), instance_count), account_id)

    connections = collections.Counter() if connections is None else connections
    now = datetime.datetime.utcnow()
    start_time = LAST_TIMES.get(account_id, now - datetime.timedelta(minutes=10))
    reader = FlowLogsReader('vpc-flowgroup', region_name=region, start_time=start_time, end_time=now)
    reader.logs_client = session.client('logs')
    record_count = 0
    new_connections = 0
    for record in reader:
        # just consider accepted packets
        if record.action == 'ACCEPT':
            record_count += 1
            src = ipaddress.ip_address(record.srcaddr)
            # only look at packets received at public interfaces
            if record.interface_id in interfaces and not src.is_private:
                name = get_name(record.srcaddr)
                dest = interfaces.get(record.interface_id, {}).get('Description')
                if not dest or dest.startswith('Primary'):
                    # EC2 instance
                    if record.srcaddr in AWS_IPS and AWS_S3_DOMAIN_PATTERN.match(name.split('/')[0]):
                        # ignore S3 public IPs
                        # (most probably packets from S3 to public EC2)
                        continue
                    dest = local_names.get(record.dstaddr, record.dstaddr)
                elif dest.startswith('ELB'):
                    # ELB
                    words = dest.split()
                    dest = lb_dns_names.get(words[-1], dest)
                elif dest.startswith('RDS'):
                    # RDS instance
                    public_ip = interfaces.get(record.interface_id, {}).get('Association', {}).get('PublicIp', '')
                    dest = NAMES.get(public_ip, 'RDS/' + public_ip)
                elif dest:
                    dest += '/' + interfaces.get(record.interface_id, {}).get('Association', {}).get('PublicIp', '')
                if 'NAT' not in dest and 'Odd' not in dest:
                    conn = (name, dest, record.dstport)
                    if conn not in connections:
                        new_connections += 1
                    connections[conn] += 1
    logging.info('%s: Got {} records and {} new connections'.format(record_count, new_connections), account_id)
    LAST_TIMES[account_id] = now
    return connections
예제 #16
0
    def test_threads(self):
        inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
            thread_count=1,
        )

        paginators = []

        def _get_paginator(operation):
            nonlocal paginators

            paginator = MagicMock()
            if operation == 'describe_log_streams':
                paginator.paginate.return_value = [
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'too_late',
                                'firstEventTimestamp': inst.end_ms,
                                'lastEventTimestamp': inst.start_ms,
                            },
                            {
                                'logStreamName': 'too_late',
                                'firstEventTimestamp': inst.end_ms - 1,
                                'lastEventTimestamp': (
                                    inst.start_ms - LAST_EVENT_DELAY_MSEC - 1
                                ),
                            },
                        ],
                    },
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'first_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                            {
                                'logStreamName': 'second_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                        ],
                    },
                ]
            elif operation == 'filter_log_events':
                paginator.paginate.return_value = [
                    {
                        'events': [
                            {'message': V2_RECORDS[0]},
                            {'message': V2_RECORDS[1]},
                        ],
                    },
                    {
                        'events': [
                            {'message': V2_RECORDS[2]},
                            {'message': V2_RECORDS[3]},
                        ],
                    },
                ]
            else:
                self.fail('invalid operation')

            paginators.append(paginator)
            return paginator

        self.mock_client.get_paginator.side_effect = _get_paginator
        events = list(inst)
        self.assertEqual(len(events), 8)

        paginators[0].paginate.assert_called_once_with(
            logGroupName='group_name',
            orderBy='LastEventTime',
            descending=True,
        )
        paginators[1].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['first_stream'],
        )
        paginators[2].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['second_stream'],
        )
예제 #17
0
class FlowLogsReaderTestCase(TestCase):
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.start_ms // 1000),
            self.start_time,
        )

        self.assertEqual(
            datetime.utcfromtimestamp(self.inst.end_ms // 1000), self.end_time
        )

        self.assertEqual(self.inst.paginator_kwargs['filterPattern'], 'REJECT')

    @patch('flowlogs_reader.flowlogs_reader.boto3.client', autospec=True)
    def test_region_name(self, mock_client):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_client.assert_called_with('logs', region_name='some-region')

        # None specified
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs')

    @patch('flowlogs_reader.flowlogs_reader.boto3.client', autospec=True)
    def test_get_fields(self, mock_client):
        cwl_client = MagicMock()
        ec2_client = mock_client.return_value
        ec2_client.describe_flow_logs.return_value = {
            'FlowLogs': [
                {'LogFormat': '${srcaddr} ${dstaddr} ${start} ${log-status}'}
            ]
        }
        reader = FlowLogsReader(
            'some_group',
            boto_client=cwl_client,
            fields=None,
        )
        self.assertEqual(
            reader.fields, ('srcaddr', 'dstaddr', 'start', 'log_status')
        )
        ec2_client.describe_flow_logs.assert_called_once_with(
            Filters=[{'Name': 'log-group-name', 'Values': ['some_group']}]
        )

    def test_read_streams(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

        actual = list(self.inst._read_streams())
        expected = []
        for page in paginator.paginate.return_value:
            expected += page['events']
        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_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_iteration_unexpecetd_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': V2_RECORDS[0]}
            yield {'events': [event_0]}
            raise PaginationError(message='other error')

        self.mock_client.get_paginator.return_value.paginate.side_effect = (
            _get_paginator
        )

        # Fail for unexpected PaginationError
        self.assertRaises(PaginationError, lambda: list(self.inst))

    def test_threads(self):
        inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
            thread_count=1,
        )

        paginators = []

        def _get_paginator(operation):
            nonlocal paginators

            paginator = MagicMock()
            if operation == 'describe_log_streams':
                paginator.paginate.return_value = [
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'too_late',
                                'firstEventTimestamp': inst.end_ms,
                                'lastEventTimestamp': inst.start_ms,
                            },
                            {
                                'logStreamName': 'too_late',
                                'firstEventTimestamp': inst.end_ms - 1,
                                'lastEventTimestamp': (
                                    inst.start_ms - LAST_EVENT_DELAY_MSEC - 1
                                ),
                            },
                        ],
                    },
                    {
                        'logStreams': [
                            {
                                'logStreamName': 'first_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                            {
                                'logStreamName': 'second_stream',
                                'firstEventTimestamp': inst.start_ms,
                                'lastEventTimestamp': inst.end_ms,
                            },
                        ],
                    },
                ]
            elif operation == 'filter_log_events':
                paginator.paginate.return_value = [
                    {
                        'events': [
                            {'message': V2_RECORDS[0]},
                            {'message': V2_RECORDS[1]},
                        ],
                    },
                    {
                        'events': [
                            {'message': V2_RECORDS[2]},
                            {'message': V2_RECORDS[3]},
                        ],
                    },
                ]
            else:
                self.fail('invalid operation')

            paginators.append(paginator)
            return paginator

        self.mock_client.get_paginator.side_effect = _get_paginator
        events = list(inst)
        self.assertEqual(len(events), 8)

        paginators[0].paginate.assert_called_once_with(
            logGroupName='group_name',
            orderBy='LastEventTime',
            descending=True,
        )
        paginators[1].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['first_stream'],
        )
        paginators[2].paginate.assert_called_once_with(
            logGroupName='group_name',
            startTime=inst.start_ms,
            endTime=inst.end_ms,
            interleaved=True,
            filterPattern='REJECT',
            logStreamNames=['second_stream'],
        )
예제 #18
0
class FlowLogsReaderTestCase(TestCase):
    @patch('flowlogs_reader.flowlogs_reader.boto3', autospec=True)
    def setUp(self, mock_boto3):
        self.mock_client = MagicMock()
        mock_boto3.client.return_value = self.mock_client

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(datetime.utcfromtimestamp(self.inst.start_ms // 1000),
                         self.start_time)

        self.assertEqual(datetime.utcfromtimestamp(self.inst.end_ms // 1000),
                         self.end_time)

    @patch('flowlogs_reader.flowlogs_reader.boto3.client', autospec=True)
    def test_region(self, mock_client):
        # Region specified
        FlowLogsReader('some_group', region_name='some-region')
        mock_client.assert_called_with('logs', region_name='some-region')

        # No region specified - assume configuration file worked
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs')

        # No region specified and no configuration file
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_client.side_effect = mock_response
        FlowLogsReader('some_group')
        mock_client.assert_called_with('logs', region_name=DEFAULT_REGION_NAME)

    def test_read_streams(self):
        response_list = [
            {
                'events': [0],
                'nextToken': 'token_0'
            },
            {
                'events': [1, 2],
                'nextToken': 'token_1'
            },
            {
                'events': [3, 4, 5],
                'nextToken': None
            },
            {
                'events': [6],
                'nextForwardToken': 'token_2'
            },  # Unreachable
        ]

        def mock_filter(*args, **kwargs):
            return response_list.pop(0)

        self.mock_client.filter_log_events.side_effect = mock_filter

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        self.assertEqual(actual, expected)

    def test_iteration(self):
        response_list = [
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[0]
                    },
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[1]
                    },
                ],
                'nextToken':
                'token_0',
            },
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[2]
                    },
                    {
                        'logStreamName': 'log_1',
                        'message': SAMPLE_RECORDS[3]
                    },
                    {
                        'logStreamName': 'log_2',
                        'message': SAMPLE_RECORDS[4]
                    },
                ],
            },
        ]

        def mock_filter(*args, **kwargs):
            return response_list.pop(0)

        self.mock_client.filter_log_events.side_effect = mock_filter

        # Calling list on the instance causes it to iterate through all records
        actual = list(self.inst)
        expected = [FlowRecord.from_message(x) for x in SAMPLE_RECORDS]
        self.assertEqual(actual, expected)
예제 #19
0
class FlowLogsReaderTestCase(TestCase):
    def setUp(self):
        self.mock_client = MagicMock()

        self.start_time = datetime(2015, 8, 12, 12, 0, 0)
        self.end_time = datetime(2015, 8, 12, 13, 0, 0)

        self.inst = FlowLogsReader(
            'group_name',
            start_time=self.start_time,
            end_time=self.end_time,
            filter_pattern='REJECT',
            boto_client=self.mock_client,
        )

    def test_init(self):
        self.assertEqual(self.inst.log_group_name, 'group_name')

        self.assertEqual(datetime.utcfromtimestamp(self.inst.start_ms // 1000),
                         self.start_time)

        self.assertEqual(datetime.utcfromtimestamp(self.inst.end_ms // 1000),
                         self.end_time)

        self.assertEqual(self.inst.paginator_kwargs['filterPattern'], 'REJECT')

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_region_name(self, mock_session):
        # Region specified for session
        FlowLogsReader('some_group', region_name='some-region')
        mock_session.Session.assert_called_with(region_name='some-region')

        # Region specified for client, not for session
        FlowLogsReader('some_group',
                       boto_client_kwargs={'region_name': 'my-region'})
        mock_session.Session().client.assert_called_with(
            'logs', region_name='my-region')

        # No region specified for session or client - use the default
        def mock_response(*args, **kwargs):
            if 'region_name' not in kwargs:
                raise NoRegionError

        mock_session.Session().client.side_effect = mock_response

        FlowLogsReader('some_group')
        mock_session.Session().client.assert_called_with(
            'logs', region_name=DEFAULT_REGION_NAME)

    @patch('flowlogs_reader.flowlogs_reader.boto3.session', autospec=True)
    def test_profile_name(self, mock_session):
        # profile_name specified
        FlowLogsReader('some_group', profile_name='my-profile')
        mock_session.Session.assert_called_with(profile_name='my-profile')

        # No profile specified
        FlowLogsReader('some_group')
        mock_session.Session.assert_called_with()

    def test_read_streams(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [0]
            },
            {
                'events': [1, 2]
            },
            {
                'events': [3, 4, 5]
            },
        ]

        self.mock_client.get_paginator.return_value = paginator

        actual = list(self.inst._read_streams())
        expected = [0, 1, 2, 3, 4, 5]
        self.assertEqual(actual, expected)

    def test_iteration(self):
        paginator = MagicMock()
        paginator.paginate.return_value = [
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[0]
                    },
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[1]
                    },
                ],
            },
            {
                'events': [
                    {
                        'logStreamName': 'log_0',
                        'message': SAMPLE_RECORDS[2]
                    },
                    {
                        'logStreamName': 'log_1',
                        'message': SAMPLE_RECORDS[3]
                    },
                    {
                        'logStreamName': 'log_2',
                        'message': SAMPLE_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_message(x) for x in SAMPLE_RECORDS]
        self.assertEqual(actual, expected)

    def test_iteration_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]}
            event_1 = {'logStreamName': 'log_0', 'message': SAMPLE_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)
        expected = [FlowRecord.from_message(x) for x in SAMPLE_RECORDS[:2]]
        self.assertEqual(actual, expected)

    def test_iteration_unexpecetd_error(self):
        # Simulate the paginator failing
        def _get_paginator(*args, **kwargs):
            event_0 = {'logStreamName': 'log_0', 'message': SAMPLE_RECORDS[0]}
            yield {'events': [event_0]}
            raise PaginationError(message='other error')

        self.mock_client.get_paginator.return_value.paginate.side_effect = (
            _get_paginator)

        # Fail for unexpected PaginationError
        self.assertRaises(PaginationError, lambda: list(self.inst))