From 09e4c2b03a0288c68c8555783462a88c54dd4e85 Mon Sep 17 00:00:00 2001 From: joon Date: Wed, 17 Apr 2024 02:52:54 -0700 Subject: [PATCH] AbuseIPDB tests and tweaks (#74) --- tests/conftest.py | 13 ++- tests/test_cli.py | 46 ++++++++- tests/test_clients.py | 29 ++++-- tests/test_data/abuseipdb_1.1.1.1_green.json | 1 + tests/test_data/abuseipdb_1.1.1.1_raw.json | 1 + tests/test_data/abuseipdb_1.1.1.1_red.json | 1 + tests/test_data/abuseipdb_1.1.1.1_yellow.json | 1 + tests/test_data/abuseipdb_one.json | 1 + tests/test_handlers.py | 20 ++++ tests/test_ui_domain_view.py | 24 ++++- tests/test_ui_ip_view.py | 97 ++++++++++++++++++- wtfis/clients/abuseipdb.py | 22 ++--- wtfis/handlers/ip.py | 6 +- wtfis/models/abuseipdb.py | 5 +- wtfis/ui/base.py | 18 ++-- wtfis/ui/view.py | 2 +- 16 files changed, 242 insertions(+), 45 deletions(-) create mode 100644 tests/test_data/abuseipdb_1.1.1.1_green.json create mode 100644 tests/test_data/abuseipdb_1.1.1.1_raw.json create mode 100644 tests/test_data/abuseipdb_1.1.1.1_red.json create mode 100644 tests/test_data/abuseipdb_1.1.1.1_yellow.json create mode 100644 tests/test_data/abuseipdb_one.json diff --git a/tests/conftest.py b/tests/conftest.py index 83706b0..f9d382e 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 50e2fc4..1f34cd2 100644 --- a/tests/test_cli.py +++ b/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 """ diff --git a/tests/test_clients.py b/tests/test_clients.py index 871e128..6078f4b 100644 --- a/tests/test_clients.py +++ b/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" diff --git a/tests/test_data/abuseipdb_1.1.1.1_green.json b/tests/test_data/abuseipdb_1.1.1.1_green.json new file mode 100644 index 0000000..c34f760 --- /dev/null +++ b/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"}} diff --git a/tests/test_data/abuseipdb_1.1.1.1_raw.json b/tests/test_data/abuseipdb_1.1.1.1_raw.json new file mode 100644 index 0000000..2e34f54 --- /dev/null +++ b/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"}} diff --git a/tests/test_data/abuseipdb_1.1.1.1_red.json b/tests/test_data/abuseipdb_1.1.1.1_red.json new file mode 100644 index 0000000..76abc7f --- /dev/null +++ b/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"}} diff --git a/tests/test_data/abuseipdb_1.1.1.1_yellow.json b/tests/test_data/abuseipdb_1.1.1.1_yellow.json new file mode 100644 index 0000000..44de3c5 --- /dev/null +++ b/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"}} diff --git a/tests/test_data/abuseipdb_one.json b/tests/test_data/abuseipdb_one.json new file mode 100644 index 0000000..da92694 --- /dev/null +++ b/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"}} diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 0af4d81..a408ec7 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -378,6 +378,7 @@ 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() @@ -385,6 +386,7 @@ def test_fetch_data(self, ip_handler): 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:") diff --git a/tests/test_ui_domain_view.py b/tests/test_ui_domain_view.py index 487b01f..42395bf 100644 --- a/tests/test_ui_domain_view.py +++ b/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,6 +190,11 @@ 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(), @@ -196,7 +202,7 @@ def view09(test_data, mock_shodan_get_ip, mock_greynoise_get): 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 diff --git a/tests/test_ui_ip_view.py b/tests/test_ui_ip_view.py index 2bbf5b2..2945668 100644 --- a/tests/test_ui_ip_view.py +++ b/tests/test_ui_ip_view.py @@ -24,9 +24,15 @@ @pytest.fixture() -def view01(test_data, mock_ipwhois_get, mock_greynoise_get, mock_urlhaus_get): +def view01(test_data, mock_abuseipdb_get, mock_ipwhois_get, mock_greynoise_get, mock_urlhaus_get): """ 1.1.1.1 with PT whois. Complete test of all panels. Also test print(). """ ip = "1.1.1.1" + + abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_red.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(ip) + ipwhois_pool = json.loads(test_data("ipwhois_1.1.1.1.json")) ipwhois_client = IpWhoisClient() ipwhois_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, ipwhois_pool)) @@ -40,7 +46,7 @@ def view01(test_data, mock_ipwhois_get, mock_greynoise_get, mock_urlhaus_get): urlhaus_pool = json.loads(test_data("urlhaus_1.1.1.1.json")) urlhaus_client = UrlHausClient() urlhaus_client._get_host = MagicMock(side_effect=lambda ip: mock_urlhaus_get(ip, urlhaus_pool)) - urlhaus_enrich = urlhaus_client.enrich_ips("1.1.1.1") + urlhaus_enrich = urlhaus_client.enrich_ips(ip) return IpAddressView( console=Console(), @@ -48,7 +54,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(), + abuseipdb=abuseipdb_enrich, urlhaus=urlhaus_enrich, ) @@ -149,6 +155,46 @@ def view06(test_data, mock_greynoise_get): ) +@pytest.fixture() +def view07(test_data, mock_abuseipdb_get): + """ 1.1.1.1 with green AbuseIPDB score. Test AbuseIPDB only. """ + ip = "1.1.1.1" + abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_green.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(ip) + + return IpAddressView( + console=Console(), + entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))), + whois=MagicMock(), + ip_enrich=IpWhoisMap.model_validate({}), + greynoise=MagicMock(), + abuseipdb=abuseipdb_enrich, + urlhaus=MagicMock(), + ) + + +@pytest.fixture() +def view08(test_data, mock_abuseipdb_get): + """ 1.1.1.1 with yellow AbuseIPDB score. Test AbuseIPDB only. """ + ip = "1.1.1.1" + abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_yellow.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(ip) + + return IpAddressView( + console=Console(), + entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))), + whois=MagicMock(), + ip_enrich=IpWhoisMap.model_validate({}), + greynoise=MagicMock(), + abuseipdb=abuseipdb_enrich, + urlhaus=MagicMock(), + ) + + class TestView01: def test_ip_panel(self, view01, theme, display_timestamp): ip = view01.ip_panel() @@ -313,6 +359,10 @@ def test_ip_panel(self, view01, theme, display_timestamp): "GreyNoise:", spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.1.1.1")] ), + Text( + "AbuseIPDB:", + spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.1.1.1")] + ) ] assert table.columns[1].style == theme.table_value assert table.columns[1].justify == "left" @@ -328,6 +378,14 @@ def test_ip_panel(self, view01, 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), + ] + ), ] def test_whois_panel(self, view01, theme): @@ -734,3 +792,36 @@ def test_ip_panel_greynoise_only(self, view06, theme): Span(19, 26, theme.tags), ] ) + + +class TestAbuseIpDbOnly: + def test_abuseipdb_green(self, view07, theme): + ip = view07.ip_panel() + + other_section = ip.renderable.renderables[2] + + # Table + table = other_section.renderables[1] + assert table.columns[1]._cells[-1] == Text( + "0 confidence score", + spans=[ + Span(0, 1, theme.info), + Span(1, 18, "green"), + ] + ) + + def test_abuseipdb_yellow(self, view08, theme): + ip = view08.ip_panel() + + other_section = ip.renderable.renderables[2] + + # Table + table = other_section.renderables[1] + assert table.columns[1]._cells[-1] == Text( + "30 confidence score (567 reports)", + spans=[ + Span(0, 2, theme.warn), + Span(2, 19, "yellow"), + Span(19, 33, theme.table_value), + ] + ) diff --git a/wtfis/clients/abuseipdb.py b/wtfis/clients/abuseipdb.py index effd069..4fcce77 100644 --- a/wtfis/clients/abuseipdb.py +++ b/wtfis/clients/abuseipdb.py @@ -1,7 +1,5 @@ from typing import Optional -from requests.exceptions import HTTPError - from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient from wtfis.models.abuseipdb import AbuseIpDb, AbuseIpDbMap @@ -10,7 +8,7 @@ class AbuseIpDbClient(BaseRequestsClient, BaseIpEnricherClient): """ AbuseIPDB client """ - baseurl = "https://api.abuseipdb.com/api/v2/check" + baseurl = "https://api.abuseipdb.com/api/v2" def __init__(self, api_key: str) -> None: super().__init__() @@ -21,18 +19,12 @@ 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 + params = {"ipAddress": ip, "maxAgeInDays": "90"} + headers = {"key": self.api_key, "Accept": "application/json"} + + response = self._get(request="/check", headers=headers, params=params) + + return AbuseIpDb.model_validate(response["data"]) def enrich_ips(self, *ips: str) -> AbuseIpDbMap: abuseipdb_map = {} diff --git a/wtfis/handlers/ip.py b/wtfis/handlers/ip.py index 39841ff..088268c 100644 --- a/wtfis/handlers/ip.py +++ b/wtfis/handlers/ip.py @@ -43,10 +43,10 @@ def fetch_data(self): self.progress.update(task_g, completed=100) if self._abuseipdb: - task_g = self.progress.add_task(f"Fetching IP enrichments from {self._abuseipdb.name}") - self.progress.update(task_g, advance=50) + task_a = self.progress.add_task(f"Fetching IP enrichments from {self._abuseipdb.name}") + self.progress.update(task_a, advance=50) self._fetch_abuseipdb(self.entity) - self.progress.update(task_g, completed=100) + self.progress.update(task_a, completed=100) task_w = self.progress.add_task(f"Fetching IP whois from {self._whois.name}") self.progress.update(task_w, advance=50) diff --git a/wtfis/models/abuseipdb.py b/wtfis/models/abuseipdb.py index ad78173..d79036e 100644 --- a/wtfis/models/abuseipdb.py +++ b/wtfis/models/abuseipdb.py @@ -8,11 +8,10 @@ class AbuseIpDb(BaseModel): ip_address: str = Field(alias="ipAddress") is_public: Optional[bool] = Field(None, alias="isPublic") - ip_version: Optional[int] = Field(None, alias="numDistinctUsers") - is_whitelisted: Optional[bool] = Field(None, alias="ipVersion") + ip_version: Optional[int] = Field(None, alias="ipVersion") + is_whitelisted: Optional[bool] = Field(None, alias="isWhitelisted") abuse_confidence_score: int = Field(alias="abuseConfidenceScore") country_code: Optional[str] = Field(None, alias="countryCode") - country_name: Optional[str] = Field(None, alias="countryName") usage_type: Optional[str] = Field(None, alias="usageType") isp: str domain: Optional[str] = None diff --git a/wtfis/ui/base.py b/wtfis/ui/base.py index d445ad3..98bdc43 100644 --- a/wtfis/ui/base.py +++ b/wtfis/ui/base.py @@ -265,18 +265,20 @@ def _gen_abuseipdb_tuple(self, ip: AbuseIpDb) -> Tuple[Text, Text]: # Content # - text = Text() - - score_message = f"{str(ip.abuse_confidence_score)} abuse confidence score" - abuseipdb_text: Text if ip.abuse_confidence_score == 0: - abuseipdb_text = Text(score_message, style=self.theme.info) + style = self.theme.info elif ip.abuse_confidence_score <= 30: - abuseipdb_text = Text(score_message, style=self.theme.warn) + style = self.theme.warn else: - abuseipdb_text = Text(score_message, style=self.theme.error) + style = self.theme.error + + text = Text() + (text + .append(Text(str(ip.abuse_confidence_score), style=style)) + .append(" confidence score", style=style.replace("bold ", ""))) - text.append(abuseipdb_text) + if ip.abuse_confidence_score > 0: + text.append(f" ({ip.total_reports} reports)", style=self.theme.table_value) return title, text diff --git a/wtfis/ui/view.py b/wtfis/ui/view.py index e1811c1..d34653f 100644 --- a/wtfis/ui/view.py +++ b/wtfis/ui/view.py @@ -124,7 +124,7 @@ def resolutions_panel(self) -> Optional[Panel]: if greynoise: data += [self._gen_greynoise_tuple(greynoise)] - # abuseIPDB + # AbuseIPDB abuseipdb = self._get_abuseipdb_enrichment(attributes.ip_address) if abuseipdb: data += [self._gen_abuseipdb_tuple(abuseipdb)]