diff --git a/tests/test_data/greynoise_1.1.1.1_malicious.json b/tests/test_data/greynoise_1.1.1.1_malicious.json new file mode 100644 index 0000000..ebc6350 --- /dev/null +++ b/tests/test_data/greynoise_1.1.1.1_malicious.json @@ -0,0 +1 @@ +{"1.1.1.1": {"ip": "1.1.1.1", "noise": true, "riot": false, "classification": "malicious", "name": "Cloudflare Public DNS", "link": "https://viz.greynoise.io/riot/1.1.1.1", "last_seen": "2023-02-18", "message": "Success"}} \ No newline at end of file diff --git a/tests/test_data/greynoise_1.1.1.1_unknown.json b/tests/test_data/greynoise_1.1.1.1_unknown.json new file mode 100644 index 0000000..96441a0 --- /dev/null +++ b/tests/test_data/greynoise_1.1.1.1_unknown.json @@ -0,0 +1 @@ +{"1.1.1.1": {"ip": "1.1.1.1", "noise": false, "riot": true, "classification": "unknown", "name": "Cloudflare Public DNS", "link": "https://viz.greynoise.io/riot/1.1.1.1", "last_seen": "2023-02-18", "message": "Success"}} \ No newline at end of file diff --git a/tests/test_ui_ip_view.py b/tests/test_ui_ip_view.py index b9a48a9..c0e9a55 100644 --- a/tests/test_ui_ip_view.py +++ b/tests/test_ui_ip_view.py @@ -45,7 +45,7 @@ def view01(test_data, mock_ipwhois_get, mock_greynoise_get): @pytest.fixture() def view02(test_data, mock_shodan_get_ip, mock_greynoise_get): - """ 1.1.1.1 with Shodan. Test the whole IP panel. """ + """ 1.1.1.1 with Shodan and Greynoise. Test the whole IP panel. """ ip = "1.1.1.1" shodan_pool = json.loads(test_data("shodan_1.1.1.1.json")) shodan_client = ShodanClient(MagicMock()) @@ -93,6 +93,42 @@ def view04(test_data): ) +@pytest.fixture() +def view05(test_data, mock_greynoise_get): + """ 1.1.1.1 with alt Greynoise results. Test Greynoise only. """ + ip = "1.1.1.1" + greynoise_pool = json.loads(test_data("greynoise_1.1.1.1_malicious.json")) + greynoise_client = GreynoiseClient("dummykey") + greynoise_client.get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_enrich = greynoise_client.single_get_ip(ip) + + return IpAddressView( + console=Console(), + entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))), + whois=MagicMock(), + ip_enrich=IpWhoisMap(__root__={}), + greynoise=greynoise_enrich, + ) + + +@pytest.fixture() +def view06(test_data, mock_greynoise_get): + """ 1.1.1.1 with another alt Greynoise result (unknown class). Test Greynoise only. """ + ip = "1.1.1.1" + greynoise_pool = json.loads(test_data("greynoise_1.1.1.1_unknown.json")) + greynoise_client = GreynoiseClient("dummykey") + greynoise_client.get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_enrich = greynoise_client.single_get_ip(ip) + + return IpAddressView( + console=Console(), + entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))), + whois=MagicMock(), + ip_enrich=IpWhoisMap(__root__={}), + greynoise=greynoise_enrich, + ) + + class TestView01: def test_ip_panel(self, view01, theme, display_timestamp): ip = view01.ip_panel() @@ -452,3 +488,40 @@ def test_ip_panel(self, view04, display_timestamp): Text("0"), display_timestamp("2022-09-03T16:58:45Z"), ] + + +class TestView05: + def test_ip_panel_greynoise_only(self, view05, theme): + ip = view05.ip_panel() + + # Table + table = ip.renderable.renderables[1] + assert table.columns[1]._cells[-1] == Text( + "✗ riot ✓ noise ! malicious", + spans=[ + Span(0, 1, theme.warn), + Span(2, 6, theme.tags), + Span(8, 9, theme.info), + Span(10, 15, theme.tags), + Span(17, 18, theme.error), + Span(19, 28, theme.tags_red), + ] + ) + + +class TestView06: + def test_ip_panel_greynoise_only(self, view06, theme): + ip = view06.ip_panel() + + # Table + table = ip.renderable.renderables[1] + assert table.columns[1]._cells[-1] == Text( + "✓ riot ✗ noise ? unknown", + spans=[ + Span(0, 1, theme.info), + Span(2, 6, theme.tags), + Span(8, 9, theme.warn), + Span(10, 15, theme.tags), + Span(19, 26, theme.tags), + ] + ) diff --git a/wtfis/handlers/base.py b/wtfis/handlers/base.py index 3a39c02..c9eeda7 100644 --- a/wtfis/handlers/base.py +++ b/wtfis/handlers/base.py @@ -63,7 +63,7 @@ def __init__( self.vt_info: Union[Domain, IpAddress] = None # type: ignore self.ip_enrich: Union[IpWhoisMap, ShodanIpMap] = None # type: ignore self.whois: WhoisType = None # type: ignore - self.greynoise: GreynoiseIpMap = None # type: ignore + self.greynoise: GreynoiseIpMap = GreynoiseIpMap.empty() # Warning messages container self.warnings: List[str] = [] diff --git a/wtfis/handlers/domain.py b/wtfis/handlers/domain.py index d5e4fdd..ab5423f 100644 --- a/wtfis/handlers/domain.py +++ b/wtfis/handlers/domain.py @@ -64,11 +64,11 @@ def _fetch_greynoise(self) -> None: except HTTPError as e: # With warning message if e.response.status_code in (400, 429, 500): - self.greynoise = GreynoiseIpMap(__root__={}) + self.greynoise = GreynoiseIpMap.empty() self.warnings.append(f"Could not fetch Greynoise: {e}") # No warning message elif e.response.status_code == 404: - self.greynoise = GreynoiseIpMap(__root__={}) + self.greynoise = GreynoiseIpMap.empty() else: raise diff --git a/wtfis/handlers/ip.py b/wtfis/handlers/ip.py index be5c4e6..e07c728 100644 --- a/wtfis/handlers/ip.py +++ b/wtfis/handlers/ip.py @@ -25,11 +25,11 @@ def _fetch_greynoise(self) -> None: except HTTPError as e: # With warning message if e.response.status_code in (400, 429, 500): - self.greynoise = GreynoiseIpMap(__root__={}) + self.greynoise = GreynoiseIpMap.empty() self.warnings.append(f"Could not fetch Greynoise: {e}") # No warning message elif e.response.status_code == 404: - self.greynoise = GreynoiseIpMap(__root__={}) + self.greynoise = GreynoiseIpMap.empty() else: raise diff --git a/wtfis/models/greynoise.py b/wtfis/models/greynoise.py index 3cd5af0..f1df7f8 100644 --- a/wtfis/models/greynoise.py +++ b/wtfis/models/greynoise.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from pydantic import BaseModel from typing import Dict, Optional @@ -15,3 +17,7 @@ class GreynoiseIp(BaseModel): class GreynoiseIpMap(BaseModel): __root__: Dict[str, GreynoiseIp] + + @classmethod + def empty(cls) -> GreynoiseIpMap: + return cls(__root__={}) diff --git a/wtfis/ui/base.py b/wtfis/ui/base.py index 2453774..b040da8 100644 --- a/wtfis/ui/base.py +++ b/wtfis/ui/base.py @@ -225,7 +225,9 @@ def _gen_greynoise_tuple(self, ip: GreynoiseIp) -> Tuple[Text, Text]: .append(" ") .append(Text("malicious", style=self.theme.tags_red))) else: - text.append(Text(ip.classification, style=text_style)) + (text + .append("? ") + .append(Text(ip.classification, style=text_style))) return title, text