Support for abuseIPDB (#70)

* Implemented abuseIPDB client for IP addresses

README.md

@@ -192,6 +192,14 @@ Use the `-u` or `--use-urlhaus` flag to enable URLhaus enrichment for hostnames,
 
 The `Malware URLs` field name is a hyperlink (if terminal-supported) that takes you to the specific URLhaus database page for your query.
 
+### AbuseIPDB enrichment
+
+Use the `-a` or `--use-abuseipdb` flag to enable AbuseIPDB enrichment for hostnames, domains and IPs.
+
+![image](https://github.com/zbalkan/wtfis/assets/39981909/0d48cfe4-7a99-47ae-980f-47839f4f0a96)
+
+The `AbuseIPDB` field name is a hyperlink (if terminal-supported) that takes you to the specific AbuseIPDB database page for your query.
+
 ### Display options
 
 For FQDN and domain lookups, you can increase or decrease the maximum number of displayed IP resolutions with `-m NUMBER` or `--max-resolutions=NUMBER`. The upper limit is 10. If you don't need resolutions at all, set the number to `0`.

tests/test_cli.py

@@ -1,18 +1,13 @@
 import json
 import os
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
 import pytest
 from dotenv import load_dotenv
-from pathlib import Path
 from rich.console import Console
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TaskProgressColumn,
-    TimeElapsedColumn,
-    TextColumn,
-)
-from unittest.mock import patch, MagicMock
+from rich.progress import (BarColumn, Progress, SpinnerColumn,
+                           TaskProgressColumn, TextColumn, TimeElapsedColumn)
 
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ip2whois import Ip2WhoisClient
@@ -23,24 +18,19 @@
 from wtfis.clients.virustotal import VTClient
 from wtfis.handlers.domain import DomainHandler
 from wtfis.handlers.ip import IpAddressHandler
-from wtfis.main import (
-    generate_entity_handler,
-    generate_view,
-    main,
-    parse_args,
-    parse_env,
-)
+from wtfis.main import (generate_entity_handler, generate_view, main,
+                        parse_args, parse_env)
 from wtfis.models.virustotal import Domain, IpAddress
 from wtfis.ui.view import DomainView, IpAddressView
 
-
 POSSIBLE_ENV_VARS = [
     "VT_API_KEY",
     "PT_API_KEY",
     "PT_API_USER",
     "IP2WHOIS_API_KEY",
     "SHODAN_API_KEY",
     "GREYNOISE_API_KEY",
+    "ABUSEIPDB_API_KEY",
     "WTFIS_DEFAULTS",
 ]
 
@@ -84,6 +74,7 @@ def fake_load_dotenv_1(tmp_path):
         "IP2WHOIS_API_KEY": "alice",
         "SHODAN_API_KEY": "hunter2",
         "GREYNOISE_API_KEY": "upupdowndown",
+        "ABUSEIPDB_API_KEY": "dummy",
     }
     return fake_load_dotenv(tmp_path, fake_env_vars)
 
@@ -274,6 +265,7 @@ def test_env_file(self, fake_load_dotenv_1):
             assert os.environ["IP2WHOIS_API_KEY"] == "alice"
             assert os.environ["SHODAN_API_KEY"] == "hunter2"
             assert os.environ["GREYNOISE_API_KEY"] == "upupdowndown"
+            assert os.environ["ABUSEIPDB_API_KEY"] == "dummy"
         unset_env_vars()
 
     @patch("wtfis.main.load_dotenv", MagicMock())
@@ -463,6 +455,7 @@ def test_view_domain_1(self, m_domain_view, test_data):
             ip_enricher_client=MagicMock(),
             whois_client=MagicMock(),
             greynoise_client=MagicMock(),
+            abuseipdb_client=MagicMock(),
             urlhaus_client=MagicMock(),
         )
         entity.vt_info = Domain.model_validate(json.loads(test_data("vt_domain_gist.json")))
@@ -482,6 +475,7 @@ def test_view_ip_1(self, m_ip_view, test_data):
             ip_enricher_client=MagicMock(),
             whois_client=MagicMock(),
             greynoise_client=MagicMock(),
+            abuseipdb_client=MagicMock(),
             urlhaus_client=MagicMock(),
         )
         entity.vt_info = IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json")))

tests/test_clients.py

@@ -1,6 +1,7 @@
 import json
 import pytest
 from unittest.mock import MagicMock, patch
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 
 from wtfis.clients.base import requests
 from wtfis.clients.greynoise import GreynoiseClient
@@ -13,6 +14,11 @@
 from wtfis.models.ipwhois import IpWhoisMap
 
 
+@pytest.fixture()
+def abuseipdb_client():
+    return AbuseIpDbClient("dummykey")
+
+
 @pytest.fixture()
 def greynoise_client():
     return GreynoiseClient("dummykey")
@@ -96,6 +102,12 @@ 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_handlers.py

@@ -1,9 +1,11 @@
 import json
+from unittest.mock import MagicMock, patch
+
 import pytest
 from requests.exceptions import ConnectionError
 from rich.console import Console
-from unittest.mock import MagicMock, patch
 
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.base import requests
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ipwhois import IpWhoisClient
@@ -26,6 +28,7 @@ def generate_domain_handler(max_resolutions=3):
         ip_enricher_client=IpWhoisClient(),
         whois_client=PTClient("dummyuser", "dummykey"),
         greynoise_client=GreynoiseClient("dummykey"),
+        abuseipdb_client=AbuseIpDbClient("dummykey"),
         urlhaus_client=UrlHausClient(),
         max_resolutions=max_resolutions,
     )
@@ -40,6 +43,7 @@ def generate_ip_handler():
         ip_enricher_client=IpWhoisClient(),
         whois_client=PTClient("dummyuser", "dummykey"),
         greynoise_client=GreynoiseClient("dummykey"),
+        abuseipdb_client=AbuseIpDbClient("dummykey"),
         urlhaus_client=UrlHausClient(),
     )
 

tests/test_ui_domain_view.py

@@ -1,15 +1,16 @@
-import pytest
 import json
+from unittest.mock import MagicMock
 
+import pytest
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
 from rich.text import Span, Text
-from unittest.mock import MagicMock
 
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ipwhois import IpWhoisClient
 from wtfis.clients.shodan import ShodanClient
+from wtfis.models.abuseipdb import AbuseIpDbMap
 from wtfis.clients.urlhaus import UrlHausClient
 from wtfis.models.greynoise import GreynoiseIpMap
 from wtfis.models.ip2whois import Whois as Ip2Whois
@@ -41,6 +42,7 @@ def view01(test_data, mock_ipwhois_get):
         whois=PTWhois.model_validate(json.loads(test_data("pt_whois_gist.json"))),
         ip_enrich=ip_enrich,
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=UrlHausMap.model_validate({}),
     )
 
@@ -58,6 +60,7 @@ def view02(test_data):
         whois=VTWhois.model_validate(json.loads(test_data("vt_whois_gist.json"))),
         ip_enrich=IpWhoisMap.model_validate({}),
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=MagicMock(),
         max_resolutions=1,
     )
@@ -73,6 +76,7 @@ def view03(test_data):
         whois=VTWhois.model_validate(json.loads(test_data("vt_whois_bbc.json"))),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -90,6 +94,7 @@ def view04(test_data):
         whois=MagicMock(),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -104,6 +109,7 @@ def view05(test_data):
         whois=MagicMock(),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -118,6 +124,7 @@ def view06(test_data):
         whois=VTWhois.model_validate(json.loads(test_data("vt_whois_example_2.json"))),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -139,6 +146,7 @@ def view07(test_data, mock_shodan_get_ip):
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=MagicMock(),
     )
 
@@ -160,6 +168,7 @@ def view08(test_data, mock_shodan_get_ip):
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=MagicMock(),
         max_resolutions=1,
     )
@@ -187,6 +196,7 @@ def view09(test_data, mock_shodan_get_ip, mock_greynoise_get):
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=greynoise_enrich,
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
         max_resolutions=1,
     )
@@ -202,6 +212,7 @@ def view10(test_data):
         whois=VTWhois.model_validate(json.loads(test_data("vt_whois_foo.json"))),
         ip_enrich=MagicMock(),
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=MagicMock(),
     )
 
@@ -223,6 +234,7 @@ def view11(test_data, mock_shodan_get_ip):
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=MagicMock(),
     )
 
@@ -237,6 +249,7 @@ def view12(test_data):
         whois=Ip2Whois.model_validate(json.loads(test_data("ip2whois_whois_hotmail.json"))),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -251,6 +264,7 @@ def view13(test_data):
         whois=Ip2Whois.model_validate(json.loads(test_data("ip2whois_whois_bbc.json"))),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -277,6 +291,7 @@ def view14(test_data, mock_ipwhois_get, mock_urlhaus_get):
         whois=PTWhois.model_validate(json.loads(test_data("pt_whois_gist.json"))),
         ip_enrich=ip_enrich,
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=AbuseIpDbMap.model_validate({}),
         urlhaus=urlhaus_enrich,
     )
 

tests/test_ui_ip_view.py

@@ -6,6 +6,7 @@
 from rich.table import Table
 from rich.text import Span, Text
 from unittest.mock import MagicMock
+from wtfis.clients.abuseipdb import AbuseIpDbClient
 
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ipwhois import IpWhoisClient
@@ -47,6 +48,7 @@ def view01(test_data, mock_ipwhois_get, mock_greynoise_get, mock_urlhaus_get):
         whois=PTWhois.model_validate(json.loads(test_data("pt_whois_1.1.1.1.json"))),
         ip_enrich=ip_enrich,
         greynoise=greynoise_enrich,
+        abuseipdb=MagicMock(),
         urlhaus=urlhaus_enrich,
     )
 
@@ -71,6 +73,7 @@ def view02(test_data, mock_shodan_get_ip, mock_greynoise_get):
         whois=MagicMock(),
         ip_enrich=ip_enrich,
         greynoise=greynoise_enrich,
+        abuseipdb=MagicMock(),
         urlhaus=UrlHausMap.model_validate({}),
     )
 
@@ -84,6 +87,7 @@ def view03(test_data):
         whois=VTWhois.model_validate(json.loads(test_data("vt_whois_1.1.1.1.json"))),
         ip_enrich=MagicMock(),
         greynoise=MagicMock(),
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 
@@ -100,6 +104,7 @@ def view04(test_data):
         whois=MagicMock(),
         ip_enrich=IpWhoisMap.model_validate({}),
         greynoise=GreynoiseIpMap.model_validate({}),
+        abuseipdb=MagicMock(),
         urlhaus=UrlHausMap.model_validate({}),
     )
 
@@ -119,6 +124,7 @@ def view05(test_data, mock_greynoise_get):
         whois=MagicMock(),
         ip_enrich=IpWhoisMap.model_validate({}),
         greynoise=greynoise_enrich,
+        abuseipdb=MagicMock(),
         urlhaus=UrlHausMap.model_validate({}),
     )
 
@@ -138,6 +144,7 @@ def view06(test_data, mock_greynoise_get):
         whois=MagicMock(),
         ip_enrich=IpWhoisMap.model_validate({}),
         greynoise=greynoise_enrich,
+        abuseipdb=MagicMock(),
         urlhaus=MagicMock(),
     )
 

wtfis/clients/abuseipdb.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+from requests.exceptions import HTTPError
+
+from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient
+from wtfis.models.abuseipdb import AbuseIpDb, AbuseIpDbMap
+
+
+class AbuseIpDbClient(BaseRequestsClient, BaseIpEnricherClient):
+    """
+    AbuseIPDB client
+    """
+    baseurl = "https://api.abuseipdb.com/api/v2/check"
+
+    def __init__(self, api_key: str) -> None:
+        super().__init__()
+        self.api_key = api_key
+
+    @property
+    def name(self) -> str:
+        return "AbuseIPDB"
+
+    def _get_ip(self, ip: str) -> Optional[AbuseIpDb]:
+        # Let a 404 or invalid IP pass
+        try:
+            params = {"ipAddress": ip, "maxAgeInDays": "90"}
+            headers = {"key": self.api_key, "Accept": "application/json"}
+
+            response = self._get(request="", headers=headers, params=params)
+
+            return AbuseIpDb.model_validate(response["data"])
+        except HTTPError as e:
+            if e.response.status_code == 404:
+                return None
+            raise
+
+    def enrich_ips(self, *ips: str) -> AbuseIpDbMap:
+        abuseipdb_map = {}
+        for ip in ips:
+            ip_data = self._get_ip(ip)
+            if ip_data:
+                abuseipdb_map[ip_data.ip_address] = ip_data
+        return AbuseIpDbMap.model_validate(abuseipdb_map)