Refactor: Extract BaseClient from Client to share logic between clients

Author: Chocapikk

leakix/__init__.py

@@ -1,7 +1,7 @@
 from importlib.metadata import version
 
+from leakix.base import HostResult as HostResult
 from leakix.client import Client as Client
-from leakix.client import HostResult as HostResult
 from leakix.client import Scope as Scope
 from leakix.domain import L9Subdomain as L9Subdomain
 from leakix.field import (

leakix/base.py

@@ -0,0 +1,73 @@
+"""Shared logic between sync and async LeakIX clients."""
+
+import dataclasses
+from importlib.metadata import version
+from typing import Any, cast
+
+from l9format import l9format
+from l9format.l9format import Model
+
+from leakix.domain import L9Subdomain
+from leakix.plugin import APIResult
+from leakix.response import AbstractResponse
+
+DEFAULT_URL = "https://leakix.net"
+
+
+@dataclasses.dataclass
+class HostResult(Model):
+    Services: list[l9format.L9Event] | None = None
+    Leaks: list[l9format.L9Event] | None = None
+
+
+class BaseClient:
+    """Shared initialization and response transformation logic."""
+
+    MAX_RESULTS_PER_PAGE = 20
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = DEFAULT_URL,
+    ) -> None:
+        self.api_key = api_key
+        self.base_url = base_url if base_url else DEFAULT_URL
+        self.headers: dict[str, str] = {
+            "Accept": "application/json",
+            "User-agent": f"leakix-client-python/{version('leakix')}",
+        }
+        if api_key:
+            self.headers["api-key"] = api_key
+
+    @staticmethod
+    def _parse_events(response: AbstractResponse) -> AbstractResponse:
+        """Parse raw JSON dicts into L9Event objects on a success response."""
+        if response.is_success():
+            response.response_json = [
+                l9format.L9Event.from_dict(res) for res in response.response_json
+            ]
+        return response
+
+    @staticmethod
+    def _parse_host_result(response: AbstractResponse) -> AbstractResponse:
+        """Parse a host/domain response into {services, leaks} format."""
+        if response.is_success():
+            data: dict[str, Any] = response.json()
+            formatted = cast(HostResult, HostResult.from_dict(data))
+            response.response_json = {
+                "services": formatted.Services,
+                "leaks": formatted.Leaks,
+            }
+        return response
+
+    @staticmethod
+    def _parse_plugins(response: AbstractResponse) -> AbstractResponse:
+        if response.is_success():
+            response.response_json = [APIResult.from_dict(d) for d in response.json()]
+        return response
+
+    @staticmethod
+    def _parse_subdomains(response: AbstractResponse) -> AbstractResponse:
+        if response.is_success():
+            response.response_json = [L9Subdomain.from_dict(d) for d in response.json()]
+        return response

leakix/client.py

@@ -1,15 +1,12 @@
-import dataclasses
 import json
 from enum import Enum
-from importlib.metadata import version
-from typing import Any, cast
+from typing import Any
 
 import requests
 from l9format import l9format
-from l9format.l9format import Model
 
-from leakix.domain import L9Subdomain
-from leakix.plugin import APIResult
+from leakix.base import BaseClient
+from leakix.base import HostResult as HostResult
 from leakix.query import EmptyQuery, Query
 from leakix.response import (
     AbstractResponse,
@@ -24,32 +21,7 @@ class Scope(Enum):
     LEAK = "leak"
 
 
-@dataclasses.dataclass
-class HostResult(Model):
-    Services: list[l9format.L9Event] | None = None
-    Leaks: list[l9format.L9Event] | None = None
-
-
-DEFAULT_URL = "https://leakix.net"
-
-
-class Client:
-    MAX_RESULTS_PER_PAGE = 20
-
-    def __init__(
-        self,
-        api_key: str | None = None,
-        base_url: str | None = DEFAULT_URL,
-    ) -> None:
-        self.api_key = api_key
-        self.base_url = base_url if base_url else DEFAULT_URL
-        self.headers: dict[str, str] = {
-            "Accept": "application/json",
-            "User-agent": f"leakix-client-python/{version('leakix')}",
-        }
-        if api_key:
-            self.headers["api-key"] = api_key
-
+class Client(BaseClient):
     def __get(self, url: str, params: dict[str, Any] | None) -> AbstractResponse:
         r = requests.get(
             url,
@@ -101,59 +73,34 @@ def get(
         else:
             serialized_query = " ".join(q.serialize() for q in queries)
         url = f"{self.base_url}/search"
-        r = self.__get(
+        return self.__get(
             url=url,
             params={
                 "scope": scope.value,
                 "q": serialized_query,
                 "page": page,
             },
         )
-        return r
 
     def get_service(
         self, queries: list[Query] | None = None, page: int = 0
     ) -> AbstractResponse:
-        """
-        Shortcut for `get` with the scope `Scope.Service`.
-
-        """
-        r = self.get(Scope.SERVICE, queries=queries, page=page)
-        if r.is_success():
-            r.response_json = [
-                l9format.L9Event.from_dict(res) for res in r.response_json
-            ]
-        return r
+        """Shortcut for `get` with the scope `Scope.SERVICE`."""
+        return self._parse_events(self.get(Scope.SERVICE, queries=queries, page=page))
 
     def get_leak(
         self, queries: list[Query] | None = None, page: int = 0
     ) -> AbstractResponse:
-        """
-        Shortcut for `get` with the scope `Scope.Leak`.
-        """
-        r = self.get(Scope.LEAK, queries=queries, page=page)
-        if r.is_success():
-            r.response_json = [
-                l9format.L9Event.from_dict(res) for res in r.response_json
-            ]
-        return r
+        """Shortcut for `get` with the scope `Scope.LEAK`."""
+        return self._parse_events(self.get(Scope.LEAK, queries=queries, page=page))
 
     def get_host(self, ipv4: str) -> AbstractResponse:
         """
         Returns the list of services and associated leaks for a given host. Only the ipv4 format is supported at the
         moment.
         """
         url = f"{self.base_url}/host/{ipv4}"
-        r = self.__get(url, params=None)
-        if r.is_success():
-            response_json = r.json()
-            formatted_result = cast(HostResult, HostResult.from_dict(response_json))
-            response_json = {
-                "services": formatted_result.Services,
-                "leaks": formatted_result.Leaks,
-            }
-            r.response_json = response_json
-        return r
+        return self._parse_host_result(self.__get(url, params=None))
 
     def get_plugins(self) -> AbstractResponse:
         """
@@ -166,10 +113,7 @@ def get_plugins(self) -> AbstractResponse:
         For the paid plans, have a look at https://leakix.net/plans.
         """
         url = f"{self.base_url}/api/plugins"
-        r = self.__get(url, params=None)
-        if r.is_success():
-            r.response_json = [APIResult.from_dict(d) for d in r.json()]
-        return r
+        return self._parse_plugins(self.__get(url, params=None))
 
     def get_subdomains(self, domain: str) -> AbstractResponse:
         """
@@ -178,10 +122,7 @@ def get_subdomains(self, domain: str) -> AbstractResponse:
         To get back a JSON/Python dictionary, use the method `to_dict` on the individual element of the response object.
         """
         url = f"{self.base_url}/api/subdomains/{domain}"
-        r = self.__get(url, params=None)
-        if r.is_success():
-            r.response_json = [L9Subdomain.from_dict(d) for d in r.json()]
-        return r
+        return self._parse_subdomains(self.__get(url, params=None))
 
     def bulk_export(self, queries: list[Query] | None = None) -> AbstractResponse:
         url = f"{self.base_url}/bulk/search"