File size: 2,673 Bytes
7c89ed7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | 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 TestVPNAssignment:
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.vpn_assignment import enrich
assign_results = {
'hits': [
{
'_source': {
'utctimestamp': datetime.utcnow(),
'details': {
'username': 'tester@mozilla.com',
'sourceipaddress': '1.2.3.4',
}
}
}
]
}
alert = {
'summary': 'test summary',
'details': {
'something': 'original',
'sourceipaddress': '10.48.123.13',
}
}
vpn_cidrs = [
'123.11.0.0/16',
'10.48.0.0/16',
]
search_window_hours = 6
search_fn = mock_search_fn(assign_results)
enriched = enrich(alert, search_window_hours, vpn_cidrs, search_fn)
assert enriched['details']['something'] == 'original'
assert 'vpnassignment' in enriched['details']
assign = enriched['details']['vpnassignment']
assert assign.get('username') == 'tester@mozilla.com'
assert assign.get('originalip') == '1.2.3.4'
def test_not_vpn_ip(self):
from alerts.plugins.vpn_assignment import enrich
assign_results = {
'hits': [],
}
alert = {
'summary': 'test summary',
'details': {
'something': 'original',
'sourceipaddress': '10.48.123.13',
}
}
vpn_cidrs = [
'123.11.0.0/16',
# Here the IP address of the user the alert fired on is NOT in a VPN
]
search_window_hours = 6
search_fn = mock_search_fn(assign_results)
enriched = enrich(alert, search_window_hours, vpn_cidrs, search_fn)
assert enriched['details']['something'] == 'original'
assert 'vpnassignment' not in enriched
|