Merge pull request #43 from pirxthepilot/greynoise_bugfix

Greynoise bugfix

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"}}

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"}}

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),
+            ]
+        )

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] = []

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
 

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
 

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__={})

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