def test_existing_asn_cc_always_dropped_and_new_ones_added_if_possible(self, LOGGER_mock):
"""Test if already existing asn/cc are removed and new ones are (maybe) added"""
self.enricher.gi_asn.org_by_addr.side_effect = [
pygeoip.GeoIPError,
pygeoip.GeoIPError,
"AS12345"]
self.enricher.gi_cc.country_code_by_addr.side_effect = [
pygeoip.GeoIPError,
"PL",
"UK"]
data = RecordDict({
"address": [{"ip": "127.0.0.1", "cc": "JP"},
{"ip": "192.187.0.1", "cc": "US", "asn": 424242},
{"ip": "10.15.1.255", "asn": 434343}]})
data.update(self.COMMON_DATA)
expected_num_of_warnings = 4 # 2 existing `cc` + 2 existing `asn`
self.enricher.enrich(data)
self.assertEqual(data["address"], [
{"ip": "127.0.0.1"},
{"ip": "192.187.0.1", "cc": "PL"},
{"ip": "10.15.1.255", "asn": 12345, "cc": "UK"},
])
self.assertEqual(data["enriched"], ([], {
"192.187.0.1": ["cc"],
"10.15.1.255": ["asn", "cc"],
}))
self.assertEqual(
len(LOGGER_mock.warning.mock_calls),
expected_num_of_warnings)
def test_adding_asn_cc_if_possible(self):
"""Test if asn/cc are (maybe) added"""
self.enricher.gi_asn.org_by_addr.side_effect = [
pygeoip.GeoIPError,
"AS1234",
"AS123456"]
self.enricher.gi_cc.country_code_by_addr.side_effect = [
"PL",
"UK",
pygeoip.GeoIPError]
data = RecordDict({
"address": [{"ip": "127.0.0.1"},
{"ip": "192.187.0.1"},
{"ip": "10.15.1.255"}]})
data.update(self.COMMON_DATA)
self.enricher.enrich(data)
self.assertEqual(data["address"], [
{"ip": "127.0.0.1", "cc": "PL"},
{"ip": "192.187.0.1", "asn": 1234, "cc": "UK"},
{"ip": "10.15.1.255", "asn": 123456},
])
self.assertEqual(data["enriched"], ([], {
"127.0.0.1": ["cc"],
"192.187.0.1": ["asn", "cc"],
"10.15.1.255": ["asn"],
}))
def test__url_to_fqdn_or_ip__called_for_ip_url(self):
"""Test if url_to_fqdn_or_ip is called if data does not contain address and fqdn"""
data = RecordDict({"url": "http://192.168.0.1"})
data.update(self.COMMON_DATA)
self.enricher.url_to_fqdn_or_ip = mock.MagicMock(return_value="192.168.0.1")
self.enricher.enrich(data)
self.enricher.url_to_fqdn_or_ip.assert_called_with("http://192.168.0.1")
def test__fqdn_to_ip__called(self):
"""Test if fqdn_to_ip is called if data does not contain address"""
data = RecordDict({"fqdn": "cert.pl"})
data.update(self.COMMON_DATA)
self.enricher.fqdn_to_ip = mock.MagicMock()
self.enricher.enrich(data)
self.enricher.fqdn_to_ip.assert_called_with("cert.pl")
def test__delete_too_long_address__address_is_empty(self):
parsed = RecordDict()
parsed.update({'source': 'foo.bar'})
expected = RecordDict()
expected.update({'source': 'foo.bar'})
self.meth.delete_too_long_address(parsed)
self.assertEqual(parsed, expected)
def test__fqdn_to_ip__not_called(self):
"""Test if fqdn_to_ip not called if address already present"""
data = RecordDict({
"address": [{"ip": "127.0.0.1"},
{"ip": "192.187.0.1"},
{"ip": "10.15.1.255"}]})
data.update(self.COMMON_DATA)
self.enricher.fqdn_to_ip = mock.MagicMock(return_value="127.0.0.1")
self.enricher.enrich(data)
self.assertFalse(self.enricher.fqdn_to_ip.called)
def test__ip_to_cc__called(self):
"""Test if ip_to_cc was called for all ips"""
data = RecordDict({
"address": [{"ip": "127.0.0.1"},
{"ip": "192.187.0.1"},
{"ip": "10.15.1.255"}]})
data.update(self.COMMON_DATA)
self.enricher.ip_to_cc = mock.MagicMock(return_value="")
self.enricher.enrich(data)
for addr in data["address"]:
self.enricher.ip_to_cc.assert_any_call(addr["ip"])
self.assertEqual(len(data["address"]), self.enricher.ip_to_cc.call_count)
def test_routing_key_modified(self):
"""Test if routing key after enrichement is set to "enriched.*"
when publishing to output queue"""
self.enricher.publish_output = mock.MagicMock()
data = RecordDict({
"address": [{"ip": "127.0.0.1"},
{"ip": "192.187.0.1"},
{"ip": "10.15.1.255"}]})
data.update(self.COMMON_DATA)
body = data.get_ready_json()
initial_routing_key = "event.parsed.test.test-source"
properties = None
self.enricher.input_callback(initial_routing_key, body, properties)
args, kwargs = self.enricher.publish_output.call_args
self.assertIn("routing_key", kwargs)
self.assertEqual(kwargs["routing_key"], "event.enriched.test.test-source")