AbuseIPDB tests and tweaks (#74)

tests/conftest.py

@@ -5,6 +5,7 @@
 from rich.console import RenderableType
 from rich.text import Span, Text
 
+from wtfis.models.abuseipdb import AbuseIpDb
 from wtfis.models.greynoise import GreynoiseIp
 from wtfis.models.ipwhois import IpWhois
 from wtfis.models.shodan import ShodanIp
@@ -45,6 +46,11 @@ def open_test_data(fname: str) -> str:
         return f.read()
 
 
+def abuseipdb_get_ip(ip, pool) -> AbuseIpDb:
+    """ Mock replacement for AbuseIpDbClient()._get_ip() """
+    return AbuseIpDb.model_validate(pool[ip])
+
+
 def greynoise_get(ip, pool) -> GreynoiseIp:
     """ Mock replacement for GreynoiseClient().get_ip() """
     return GreynoiseIp.model_validate(pool[ip])
@@ -61,7 +67,7 @@ def shodan_get_ip(ip, pool) -> ShodanIp:
 
 
 def urlhaus_get_host(entity, pool) -> UrlHaus:
-    """ Mock replacement for Urlhaus()._get_host() """
+    """ Mock replacement for UrlHausClient()._get_host() """
     return UrlHaus.model_validate(pool[entity])
 
 
@@ -88,6 +94,11 @@ def theme():
     return TestTheme()
 
 
+@pytest.fixture(scope="module")
+def mock_abuseipdb_get():
+    return abuseipdb_get_ip
+
+
 @pytest.fixture(scope="module")
 def mock_greynoise_get():
     return greynoise_get

tests/test_cli.py

@@ -9,6 +9,7 @@
 from rich.progress import (BarColumn, Progress, SpinnerColumn,
                            TaskProgressColumn, TextColumn, TimeElapsedColumn)
 
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ip2whois import Ip2WhoisClient
 from wtfis.clients.ipwhois import IpWhoisClient
@@ -103,7 +104,8 @@ def fake_load_dotenv_4(tmp_path):
     fake_env_vars = {
         "VT_API_KEY": "foo",
         "GREYNOISE_API_KEY": "bar",
-        "WTFIS_DEFAULTS": "-g -u",
+        "ABUSEIPDB_API_KEY": "dummy",
+        "WTFIS_DEFAULTS": "-g -u -a",
     }
     return fake_load_dotenv(tmp_path, fake_env_vars)
 
@@ -254,6 +256,32 @@ def test_urlhaus(self):
             args = parse_args()
             assert args.use_urlhaus is True
 
+    def test_abuseipdb_ok(self):
+        os.environ["ABUSEIPDB_API_KEY"] = "foo"
+        with patch("sys.argv", [
+            "main",
+            "1.1.1.1",
+            "-a",
+        ]):
+            args = parse_args()
+            assert args.use_abuseipdb is True
+        del os.environ["ABUSEIPDB_API_KEY"]
+
+    def test_abuseipdb_error(self, capsys):
+        with pytest.raises(SystemExit) as e:
+            with patch("sys.argv", [
+                "main",
+                "1.1.1.1",
+                "-a",
+            ]):
+                parse_args()
+
+        capture = capsys.readouterr()
+
+        assert capture.err == "usage: main [-h]\nmain: error: ABUSEIPDB_API_KEY is not set\n"
+        assert e.type == SystemExit
+        assert e.value.code == 2
+
 
 class TestEnvs:
     def test_env_file(self, fake_load_dotenv_1):
@@ -360,6 +388,7 @@ def test_defaults_4(self, fake_load_dotenv_4):
                 assert args.use_shodan is False
                 assert args.use_greynoise is True
                 assert args.use_urlhaus is False
+                assert args.use_abuseipdb is True
         unset_env_vars()
 
 
@@ -441,6 +470,21 @@ def test_handler_ip_1(self, fake_load_dotenv_1):
         assert entity._urlhaus is None
         unset_env_vars()
 
+    @patch("sys.argv", ["main", "1[.]1[.]1[.]1", "-s", "-g", "-u", "-a"])
+    def test_handler_ip_2(self, fake_load_dotenv_1):
+        """ IP with various options """
+        with patch("wtfis.main.load_dotenv", fake_load_dotenv_1):
+            parse_env()
+            console = Console()
+            progress = simulate_progress(console),
+            entity = generate_entity_handler(parse_args(), console, progress)
+        assert isinstance(entity._enricher, ShodanClient)
+        assert isinstance(entity._whois, PTClient)
+        assert isinstance(entity._greynoise, GreynoiseClient)
+        assert isinstance(entity._urlhaus, UrlHausClient)
+        assert isinstance(entity._abuseipdb, AbuseIpDbClient)
+        unset_env_vars()
+
 
 class TestGenView:
     """ Tests for the generate_view function """

tests/test_clients.py

@@ -1,8 +1,8 @@
 import json
 import pytest
 from unittest.mock import MagicMock, patch
-from wtfis.clients.abuseipdb import AbuseIpDbClient
 
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.base import requests
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ip2whois import Ip2WhoisClient
@@ -54,6 +54,27 @@ def virustotal_client():
     return VTClient("dummykey")
 
 
+class TestAbuseIpDbClient:
+    def test_init(self, abuseipdb_client):
+        assert abuseipdb_client.name == "AbuseIPDB"
+        assert abuseipdb_client.api_key == "dummykey"
+
+    @patch.object(requests.Session, "get")
+    def test_enrich_ips(self, mock_requests_get, test_data, abuseipdb_client):
+        mock_resp = MagicMock()
+        mock_resp.status_code = 200
+        mock_resp.json.return_value = json.loads(test_data("abuseipdb_1.1.1.1_raw.json"))
+        mock_requests_get.return_value = mock_resp
+
+        abuseipdb = abuseipdb_client.enrich_ips("thisdoesntmatter").root["1.1.1.1"]
+
+        assert abuseipdb.ip_address == "1.1.1.1"
+        assert abuseipdb.abuse_confidence_score == 100
+        assert abuseipdb.usage_type == "Data Center/Web Hosting/Transit"
+        assert abuseipdb.total_reports == 567
+        assert abuseipdb.num_distinct_users == 123
+
+
 class TestIp2WhoisClient:
     def test_init(self, ip2whois_client):
         assert ip2whois_client.api_key == "dummykey"
@@ -102,12 +123,6 @@ def test_init(self, greynoise_client):
         assert greynoise_client.api_key == "dummykey"
 
 
-class TestAbuseIPDBClient:
-    def test_init(self, abuseipdb_client):
-        assert abuseipdb_client.name == "AbuseIPDB"
-        assert abuseipdb_client.api_key == "dummykey"
-
-
 class TestIpWhoisClient:
     def test_init(self, ipwhois_client):
         assert ipwhois_client.name == "IPWhois"

tests/test_data/abuseipdb_1.1.1.1_green.json

@@ -0,0 +1 @@
+{"1.1.1.1":{"ipAddress":"1.1.1.1","isPublic":true,"ipVersion":4,"isWhitelisted":false,"abuseConfidenceScore":0,"countryCode":"US","usageType":"Data Center/Web Hosting/Transit","isp":"Acme, Inc.","domain":"example.com","hostnames":[],"isTor":false,"totalReports":567,"numDistinctUsers":123,"lastReportedAt":"2018-12-20T20:55:14+00:00"}}

tests/test_data/abuseipdb_1.1.1.1_raw.json

@@ -0,0 +1 @@
+{"data":{"ipAddress":"1.1.1.1","isPublic":true,"ipVersion":4,"isWhitelisted":false,"abuseConfidenceScore":100,"countryCode":"US","usageType":"Data Center/Web Hosting/Transit","isp":"Acme, Inc.","domain":"example.com","hostnames":[],"isTor":false,"totalReports":567,"numDistinctUsers":123,"lastReportedAt":"2018-12-20T20:55:14+00:00"}}

tests/test_data/abuseipdb_1.1.1.1_red.json

@@ -0,0 +1 @@
+{"1.1.1.1":{"ipAddress":"1.1.1.1","isPublic":true,"ipVersion":4,"isWhitelisted":false,"abuseConfidenceScore":100,"countryCode":"US","usageType":"Data Center/Web Hosting/Transit","isp":"Acme, Inc.","domain":"example.com","hostnames":[],"isTor":false,"totalReports":567,"numDistinctUsers":123,"lastReportedAt":"2018-12-20T20:55:14+00:00"}}

tests/test_data/abuseipdb_1.1.1.1_yellow.json

@@ -0,0 +1 @@
+{"1.1.1.1":{"ipAddress":"1.1.1.1","isPublic":true,"ipVersion":4,"isWhitelisted":false,"abuseConfidenceScore":30,"countryCode":"US","usageType":"Data Center/Web Hosting/Transit","isp":"Acme, Inc.","domain":"example.com","hostnames":[],"isTor":false,"totalReports":567,"numDistinctUsers":123,"lastReportedAt":"2018-12-20T20:55:14+00:00"}}

tests/test_data/abuseipdb_one.json

@@ -0,0 +1 @@
+{"1.0.0.1":{"ipAddress":"1.0.0.1","isPublic":true,"ipVersion":4,"isWhitelisted":false,"abuseConfidenceScore":100,"countryCode":"US","usageType":"Data Center/Web Hosting/Transit","isp":"Acme, Inc.","domain":"example.com","hostnames":[],"isTor":false,"totalReports":567,"numDistinctUsers":123,"lastReportedAt":"2018-12-20T20:55:14+00:00"}}

tests/test_handlers.py

@@ -378,13 +378,15 @@ def test_fetch_data(self, ip_handler):
         handler._fetch_whois = MagicMock()
         handler._fetch_greynoise = MagicMock()
         handler._fetch_urlhaus = MagicMock()
+        handler._fetch_abuseipdb = MagicMock()
 
         handler.fetch_data()
         handler._fetch_vt_ip_address.assert_called_once()
         handler._fetch_ip_enrichments.assert_called_once()
         handler._fetch_whois.assert_called_once()
         handler._fetch_greynoise.assert_called_once()
         handler._fetch_urlhaus.assert_called_once()
+        handler._fetch_abuseipdb.assert_called_once()
 
     @patch.object(requests.Session, "get")
     def test_vt_http_error(self, mock_requests_get, ip_handler, capsys):
@@ -602,3 +604,21 @@ def test_urlhaus_http_error(self, mock_requests_post, ip_handler, capsys):
         handler.print_warnings()
         capture = capsys.readouterr()
         assert capture.out.startswith("WARN: Could not fetch URLhaus: Foo bar message")
+
+    @patch.object(requests.Session, "get")
+    def test_abuseipdb_429_error(self, mock_requests_get, ip_handler, capsys):
+        """
+        Test fail open behavior of AbuseIPDB when rate limited
+        """
+        handler = ip_handler()
+        mock_resp = requests.models.Response()
+
+        mock_resp.status_code = 429
+        mock_requests_get.return_value = mock_resp
+
+        handler._fetch_abuseipdb("1.2.3.4")
+        assert handler.warnings[0].startswith("Could not fetch AbuseIPDB: 429 Client Error:")
+
+        handler.print_warnings()
+        capture = capsys.readouterr()
+        assert capture.out.startswith("WARN: Could not fetch AbuseIPDB: 429 Client Error:")

tests/test_ui_domain_view.py

@@ -7,6 +7,7 @@
 from rich.table import Table
 from rich.text import Span, Text
 
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ipwhois import IpWhoisClient
 from wtfis.clients.shodan import ShodanClient
@@ -175,8 +176,8 @@ def view08(test_data, mock_shodan_get_ip):
 
 
 @pytest.fixture()
-def view09(test_data, mock_shodan_get_ip, mock_greynoise_get):
-    """ one.one.one.one with Shodan and Greynoise. Only test resolution and IP enrich. """
+def view09(test_data, mock_shodan_get_ip, mock_greynoise_get, mock_abuseipdb_get):
+    """ one.one.one.one with Shodan, Greynoise and AbuseIPDB. Only test resolution and IP enrich. """
     resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_one.json")))
 
     shodan_pool = json.loads(test_data("shodan_one.json"))
@@ -189,14 +190,19 @@ def view09(test_data, mock_shodan_get_ip, mock_greynoise_get):
     greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool))
     greynoise_enrich = greynoise_client.enrich_ips(*resolutions.ip_list(1))
 
+    abuseipdb_pool = json.loads(test_data("abuseipdb_one.json"))
+    abuseipdb_client = AbuseIpDbClient("dummykey")
+    abuseipdb_client._get_ip = MagicMock(side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool))
+    abuseipdb_enrich = abuseipdb_client.enrich_ips(*resolutions.ip_list(1))
+
     return DomainView(
         console=Console(),
         entity=MagicMock(),
         resolutions=resolutions,
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=greynoise_enrich,
-        abuseipdb=MagicMock(),
+        abuseipdb=abuseipdb_enrich,
         urlhaus=MagicMock(),
         max_resolutions=1,
     )
@@ -1274,6 +1280,10 @@ def test_resolutions_panel(self, view09, theme, display_timestamp):
                 "GreyNoise:",
                 spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.0.0.1")]
             ),
+            Text(
+                "AbuseIPDB:",
+                spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.0.0.1")]
+            ),
         ]
         assert table.columns[1].style == theme.table_value
         assert table.columns[1].justify == "left"
@@ -1349,6 +1359,14 @@ def test_resolutions_panel(self, view09, theme, display_timestamp):
                     Span(19, 25, theme.tags_green),
                 ]
             ),
+            Text(
+                "100 confidence score (567 reports)",
+                spans=[
+                    Span(0, 3, theme.error),
+                    Span(3, 20, "red"),
+                    Span(20, 34, theme.table_value),
+                ]
+            ),
         ]
 
         # Old timestamp warning