Decouple geoip and ASN lookups from Shodan results

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
     "Topic :: Security",
 ]
 dependencies = [
-    "pydantic~=2.7.0",
+    "pydantic~=2.7.2",
     "python-dotenv~=1.0.1",
     "requests~=2.31.0",
     "rich~=13.7.1",

wtfis/clients/base.py

@@ -4,8 +4,8 @@
 
 from typing import Optional, Union
 
-from wtfis.models.common import WhoisBase
-from wtfis.types import DomainEnrichmentType, IpEnrichmentType
+from wtfis.models.base import WhoisBase
+from wtfis.models.types import DomainEnrichmentType, IpEnrichmentType
 
 
 class AbstractAttribute:

wtfis/clients/ipwhois.py

@@ -1,5 +1,6 @@
 from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient
-from wtfis.models.ipwhois import IpWhois, IpWhoisMap
+from wtfis.models.ipwhois import IpWhoisMap
+from wtfis.models.ipwhois import IpWhois
 
 from typing import Optional
 
@@ -19,9 +20,9 @@ def _get_ipwhois(self, ip: str) -> Optional[IpWhois]:
         return IpWhois.model_validate(result) if result.get("success") is True else None
 
     def enrich_ips(self, *ips: str) -> IpWhoisMap:
-        ipwhois_map = {}
+        map_ = {}
         for ip in ips:
             ipwhois = self._get_ipwhois(ip)
             if ipwhois:
-                ipwhois_map[ipwhois.ip] = ipwhois
-        return IpWhoisMap.model_validate(ipwhois_map)
+                map_[ipwhois.ip] = ipwhois
+        return IpWhoisMap.model_validate(map_)

wtfis/clients/types.py

@@ -0,0 +1,22 @@
+"""
+Type aliases
+"""
+from typing import Union
+
+from wtfis.clients.ip2whois import Ip2WhoisClient
+from wtfis.clients.ipwhois import IpWhoisClient
+from wtfis.clients.passivetotal import PTClient
+from wtfis.clients.virustotal import VTClient
+
+
+# IP geolocation and ASN client types
+IpGeoAsnClientType = Union[
+    IpWhoisClient,
+]
+
+# IP whois client types
+IpWhoisClientType = Union[
+    Ip2WhoisClient,
+    PTClient,
+    VTClient,
+]

wtfis/handlers/base.py

@@ -15,19 +15,17 @@
 
 from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.greynoise import GreynoiseClient
-from wtfis.clients.ip2whois import Ip2WhoisClient
-from wtfis.clients.ipwhois import IpWhoisClient
-from wtfis.clients.passivetotal import PTClient
 from wtfis.clients.shodan import ShodanClient
 from wtfis.clients.urlhaus import UrlHausClient
 from wtfis.clients.virustotal import VTClient
 from wtfis.models.abuseipdb import AbuseIpDbMap
-from wtfis.models.common import WhoisBase
+from wtfis.models.base import WhoisBase
 from wtfis.models.greynoise import GreynoiseIpMap
-from wtfis.models.ipwhois import IpWhoisMap
 from wtfis.models.shodan import ShodanIpMap
 from wtfis.models.urlhaus import UrlHausMap
+from wtfis.models.types import IpGeoAsnMapType
 from wtfis.models.virustotal import Domain, IpAddress
+from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType
 from wtfis.ui.theme import Theme
 from wtfis.utils import error_and_exit, refang
 
@@ -69,8 +67,9 @@ def __init__(
         console: Console,
         progress: Progress,
         vt_client: VTClient,
-        ip_enricher_client: Union[IpWhoisClient, ShodanClient],
-        whois_client: Union[Ip2WhoisClient, PTClient, VTClient],
+        ip_geoasn_client: IpGeoAsnClientType,
+        whois_client: IpWhoisClientType,
+        shodan_client: Optional[ShodanClient],
         greynoise_client: Optional[GreynoiseClient],
         abuseipdb_client: Optional[AbuseIpDbClient],
         urlhaus_client: Optional[UrlHausClient],
@@ -82,16 +81,18 @@ def __init__(
 
         # Clients
         self._vt = vt_client
-        self._enricher = ip_enricher_client
+        self._geoasn = ip_geoasn_client
         self._whois = whois_client
+        self._shodan = shodan_client
         self._greynoise = greynoise_client
         self._abuseipdb = abuseipdb_client
         self._urlhaus = urlhaus_client
 
         # Dataset containers
         self.vt_info: Union[Domain, IpAddress]
-        self.ip_enrich: Union[IpWhoisMap, ShodanIpMap] = IpWhoisMap.empty()
+        self.geoasn: IpGeoAsnMapType
         self.whois: WhoisBase
+        self.shodan: ShodanIpMap = ShodanIpMap.empty()
         self.greynoise: GreynoiseIpMap = GreynoiseIpMap.empty()
         self.abuseipdb: AbuseIpDbMap = AbuseIpDbMap.empty()
         self.urlhaus: UrlHausMap = UrlHausMap.empty()
@@ -105,9 +106,15 @@ def fetch_data(self) -> None:
         return NotImplemented  # type: ignore  # pragma: no coverage
 
     @common_exception_handler
-    @failopen_exception_handler("_enricher")
-    def _fetch_ip_enrichments(self, *ips: str) -> None:
-        self.ip_enrich = self._enricher.enrich_ips(*ips)
+    @failopen_exception_handler("_geoasn")
+    def _fetch_geoasn(self, *ips: str) -> None:
+        self.geoasn = self._geoasn.enrich_ips(*ips)
+
+    @common_exception_handler
+    @failopen_exception_handler("_shodan")
+    def _fetch_shodan(self, *ips: str) -> None:
+        if self._shodan:
+            self.shodan = self._shodan.enrich_ips(*ips)
 
     @common_exception_handler
     @failopen_exception_handler("_greynoise")

wtfis/handlers/domain.py

@@ -1,17 +1,14 @@
 """
 Logic handler for domain and hostname inputs
 """
-from typing import Optional, Union
+from typing import Optional
 
 from requests.exceptions import HTTPError
 from rich.console import Console
 from rich.progress import Progress
 
 from wtfis.clients.abuseipdb import AbuseIpDbClient
 from wtfis.clients.greynoise import GreynoiseClient
-from wtfis.clients.ip2whois import Ip2WhoisClient
-from wtfis.clients.ipwhois import IpWhoisClient
-from wtfis.clients.passivetotal import PTClient
 from wtfis.clients.shodan import ShodanClient
 from wtfis.clients.urlhaus import UrlHausClient
 from wtfis.clients.virustotal import VTClient
@@ -23,6 +20,7 @@
 )
 
 from wtfis.models.virustotal import Resolutions
+from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType
 
 
 class DomainHandler(BaseHandler):
@@ -32,15 +30,16 @@ def __init__(
         console: Console,
         progress: Progress,
         vt_client: VTClient,
-        ip_enricher_client: Union[IpWhoisClient, ShodanClient],
-        whois_client: Union[Ip2WhoisClient, PTClient, VTClient],
+        ip_geoasn_client: IpGeoAsnClientType,
+        whois_client: IpWhoisClientType,
+        shodan_client: Optional[ShodanClient],
         greynoise_client: Optional[GreynoiseClient],
         abuseipdb_client: Optional[AbuseIpDbClient],
         urlhaus_client: Optional[UrlHausClient],
         max_resolutions: int = 0,
     ):
-        super().__init__(entity, console, progress, vt_client, ip_enricher_client,
-                         whois_client, greynoise_client, abuseipdb_client, urlhaus_client)
+        super().__init__(entity, console, progress, vt_client, ip_geoasn_client, whois_client,
+                         shodan_client, greynoise_client, abuseipdb_client, urlhaus_client)
 
         # Extended attributes
         self.max_resolutions = max_resolutions
@@ -77,25 +76,31 @@ def fetch_data(self):
         self.progress.update(task_v, completed=100)
 
         if self.resolutions and self.resolutions.data:
-            task_r = self.progress.add_task(f"Fetching IP enrichments from {self._enricher.name}")
-            self.progress.update(task_r, advance=50)
-            self._fetch_ip_enrichments(*self.resolutions.ip_list(self.max_resolutions))
-            self.progress.update(task_r, completed=100)
+            task_g = self.progress.add_task(f"Fetching IP location and ASN from {self._geoasn.name}")
+            self.progress.update(task_g, advance=50)
+            self._fetch_geoasn(*self.resolutions.ip_list(self.max_resolutions))
+            self.progress.update(task_g, completed=100)
+
+            if self._shodan:
+                task_s = self.progress.add_task(f"Fetching IP data from {self._shodan.name}")
+                self.progress.update(task_s, advance=50)
+                self._fetch_shodan(*self.resolutions.ip_list(self.max_resolutions))
+                self.progress.update(task_s, completed=100)
 
             if self._greynoise:
-                task_g = self.progress.add_task(f"Fetching IP enrichments from {self._greynoise.name}")
+                task_g = self.progress.add_task(f"Fetching IP data from {self._greynoise.name}")
                 self.progress.update(task_g, advance=50)
                 self._fetch_greynoise(*self.resolutions.ip_list(self.max_resolutions))
                 self.progress.update(task_g, completed=100)
 
             if self._abuseipdb:
-                task_g = self.progress.add_task(f"Fetching IP enrichments from {self._abuseipdb.name}")
+                task_g = self.progress.add_task(f"Fetching IP data from {self._abuseipdb.name}")
                 self.progress.update(task_g, advance=50)
                 self._fetch_abuseipdb(*self.resolutions.ip_list(self.max_resolutions))
                 self.progress.update(task_g, completed=100)
 
         if self._urlhaus:
-            task_u = self.progress.add_task(f"Fetching domain enrichments from {self._urlhaus.name}")
+            task_u = self.progress.add_task(f"Fetching domain data from {self._urlhaus.name}")
             self.progress.update(task_u, advance=50)
             self._fetch_urlhaus()
             self.progress.update(task_u, completed=100)

wtfis/handlers/ip.py

@@ -25,25 +25,31 @@ def fetch_data(self):
         self._fetch_vt_ip_address()
         self.progress.update(task_v, completed=100)
 
-        task_i = self.progress.add_task(f"Fetching IP enrichments from {self._enricher.name}")
-        self.progress.update(task_i, advance=50)
-        self._fetch_ip_enrichments(self.entity)
-        self.progress.update(task_i, completed=100)
+        task_g = self.progress.add_task(f"Fetching IP location and ASN from {self._geoasn.name}")
+        self.progress.update(task_g, advance=50)
+        self._fetch_geoasn(self.entity)
+        self.progress.update(task_g, completed=100)
+
+        if self._shodan:
+            task_s = self.progress.add_task(f"Fetching IP data from {self._shodan.name}")
+            self.progress.update(task_s, advance=50)
+            self._fetch_shodan(self.entity)
+            self.progress.update(task_s, completed=100)
 
         if self._urlhaus:
-            task_u = self.progress.add_task(f"Fetching IP enrichments from {self._urlhaus.name}")
+            task_u = self.progress.add_task(f"Fetching IP data from {self._urlhaus.name}")
             self.progress.update(task_u, advance=50)
             self._fetch_urlhaus()
             self.progress.update(task_u, completed=100)
 
         if self._greynoise:
-            task_g = self.progress.add_task(f"Fetching IP enrichments from {self._greynoise.name}")
+            task_g = self.progress.add_task(f"Fetching IP data from {self._greynoise.name}")
             self.progress.update(task_g, advance=50)
             self._fetch_greynoise(self.entity)
             self.progress.update(task_g, completed=100)
 
         if self._abuseipdb:
-            task_a = self.progress.add_task(f"Fetching IP enrichments from {self._abuseipdb.name}")
+            task_a = self.progress.add_task(f"Fetching IP data from {self._abuseipdb.name}")
             self.progress.update(task_a, advance=50)
             self._fetch_abuseipdb(self.entity)
             self.progress.update(task_a, completed=100)

wtfis/main.py

@@ -8,7 +8,6 @@
 from rich.progress import Progress
 from typing import Union
 from wtfis.clients.abuseipdb import AbuseIpDbClient
-
 from wtfis.clients.greynoise import GreynoiseClient
 from wtfis.clients.ip2whois import Ip2WhoisClient
 from wtfis.clients.ipwhois import IpWhoisClient
@@ -108,12 +107,9 @@ def generate_entity_handler(
     # Virustotal client
     vt_client = VTClient(os.environ["VT_API_KEY"])
 
-    # IP enrichment client selector
-    enricher_client: Union[IpWhoisClient, ShodanClient] = (
-        ShodanClient(os.environ["SHODAN_API_KEY"])
-        if args.use_shodan
-        else IpWhoisClient()
-    )
+    # IP geolocation and ASN client selector
+    # TODO: add more options
+    ip_geoasn_client = IpWhoisClient()
 
     # Whois client selector
     # Order of use based on set envvars:
@@ -129,6 +125,12 @@ def generate_entity_handler(
     else:
         whois_client = vt_client
 
+    shodan_client = (
+        ShodanClient(os.environ["SHODAN_API_KEY"])
+        if args.use_shodan
+        else None
+    )
+
     # Greynoise client (optional)
     greynoise_client = (
         GreynoiseClient(os.environ["GREYNOISE_API_KEY"])
@@ -156,8 +158,9 @@ def generate_entity_handler(
             console=console,
             progress=progress,
             vt_client=vt_client,
-            ip_enricher_client=enricher_client,
+            ip_geoasn_client=ip_geoasn_client,
             whois_client=whois_client,
+            shodan_client=shodan_client,
             greynoise_client=greynoise_client,
             abuseipdb_client=abuseipdb_client,
             urlhaus_client=urlhaus_client,
@@ -170,8 +173,9 @@ def generate_entity_handler(
             console=console,
             progress=progress,
             vt_client=vt_client,
-            ip_enricher_client=enricher_client,
+            ip_geoasn_client=ip_geoasn_client,
             whois_client=whois_client,
+            shodan_client=shodan_client,
             greynoise_client=greynoise_client,
             abuseipdb_client=abuseipdb_client,
             urlhaus_client=urlhaus_client,
@@ -191,8 +195,9 @@ def generate_view(
             console,
             entity.vt_info,
             entity.resolutions,
+            entity.geoasn,
             entity.whois,
-            entity.ip_enrich,
+            entity.shodan,
             entity.greynoise,
             entity.abuseipdb,
             entity.urlhaus,
@@ -202,8 +207,9 @@ def generate_view(
         view = IpAddressView(
             console,
             entity.vt_info,
+            entity.geoasn,
             entity.whois,
-            entity.ip_enrich,
+            entity.shodan,
             entity.greynoise,
             entity.abuseipdb,
             entity.urlhaus,

wtfis/models/common.py → wtfis/models/base.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import abc
 import sys
 
-from pydantic import BaseModel, BeforeValidator
+from pydantic import BaseModel, BeforeValidator, ConfigDict, RootModel
 from pydantic.v1.validators import str_validator
-from typing import List, Optional
+from typing import List, Mapping, Optional
 
 if sys.version_info >= (3, 9):
     from typing import Annotated
@@ -33,3 +35,30 @@ class WhoisBase(BaseModel, abc.ABC):
     date_changed: Optional[str] = None
     date_expires: Optional[str] = None
     dnssec: Optional[str] = None
+
+
+class IpGeoAsnBase(BaseModel, abc.ABC):
+    """ Use to normalize IP geolocation and ASN fields """
+    model_config = ConfigDict(coerce_numbers_to_str=True)
+
+    ip: str
+
+    # Geolocation
+    city: Optional[str]
+    continent: Optional[str]
+    country: Optional[str]
+    region: Optional[str]
+
+    # ASN
+    asn: Optional[str]
+    org: Optional[str]
+    isp: Optional[str]
+    domain: Optional[str]
+
+
+class IpGeoAsnMapBase(RootModel, abc.ABC):
+    root: Mapping[str, IpGeoAsnBase]
+
+    @classmethod
+    def empty(cls) -> IpGeoAsnMapBase:
+        return cls.model_validate({})

wtfis/models/ip2whois.py

@@ -1,7 +1,7 @@
 from pydantic import Field, field_validator, model_validator
 from typing import List, Optional
 
-from wtfis.models.common import WhoisBase
+from wtfis.models.base import WhoisBase
 
 
 class Whois(WhoisBase):