tests/alerts/plugins/test_auth_sourceip_username.py

from datetime import datetime
import os
import sys


def mock_search_fn(results):
    '''Creates a search function that returns a set of results on each call.
    '''

    def search_fn(_query):
        return results['hits']

    return search_fn


class TestUsernameAssignment:
    def setup(self):
        self.orig_path = os.getcwd()
        self.alerts_path = os.path.join(
            os.path.dirname(__file__),
            '../../../alerts',
        )

        sys.path.insert(0, self.alerts_path)

    def teardown(self):
        os.chdir(self.orig_path)
        sys.path.remove(self.alerts_path)

        if 'lib' in sys.modules:
            del sys.modules['lib']

    def test_alert_enriched(self):
        from alerts.plugins.auth_sourceip_username import enrich

        assign_results = {
            'hits': [
                {
                    '_source': {
                        'utctimestamp': datetime.utcnow(),
                        'details': {
                            'username': 'tester@mozilla.com',
                        }
                    }
                }
            ]
        }

        alert = {
            'summary': 'test summary',
            'details': {
                'something': 'original',
                'sourceipaddress': '10.48.123.13',
            }
        }

        search_window_hours = 6

        search_fn = mock_search_fn(assign_results)

        enriched = enrich(alert, search_window_hours, search_fn)

        assign = enriched['details']

        assert assign.get('username') == 'tester@mozilla.com'