Merge pull request #1 from LeakIX/add-plugins

Add plugins

Author: dannywillems

example/example_client.py

@@ -6,7 +6,8 @@
 
 
 API_KEY = decouple.config("API_KEY")
-CLIENT = Client(api_key=API_KEY)
+BASE_URL = decouple.config("LEAKIX_HOST", default=None)
+CLIENT = Client(api_key=API_KEY, base_url=BASE_URL)
 
 
 def example_get_host_filter_plugin():
@@ -65,9 +66,17 @@ def example_get_leak_raw_query():
     )
 
 
+def example_get_plugins():
+    response = CLIENT.get_plugins()
+    for p in response.json():
+        print(p.name)
+        print(p.description)
+
+
 if __name__ == "__main__":
     example_get_host_filter_plugin()
     example_get_service_filter_plugin()
     example_get_leaks_filter_multiple_plugins()
     example_get_leaks_multiple_filter_plugins_must_not()
     example_get_leak_raw_query()
+    example_get_plugins()

leakix/__init__.py

@@ -3,5 +3,4 @@
 from leakix.plugin import *
 from leakix.query import *
 from leakix.response import *
-
-__VERSION__ = "0.1.2"
+from leakix.client import __VERSION__

leakix/client.py

@@ -2,15 +2,16 @@
 import os
 from typing import Optional, List
 import requests
-from pprint import pprint
 from l9format import l9format
 from enum import Enum
 from serde import Model, fields
 from leakix.response import SuccessResponse, ErrorResponse, RateLimitResponse
 from leakix.query import *
 from leakix.plugin import *
 from leakix.field import *
-from leakix import __VERSION__
+
+
+__VERSION__ = "0.1.2"
 
 
 class Scope(Enum):
@@ -24,28 +25,24 @@ class HostResult(Model):
 
 
 class Client:
-    BASE_URL = "https://leakix.net"
-
-    def __init__(self, api_key: Optional[str] = None):
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = "https://leakix.net",
+    ):
         self.api_key = api_key
+        self.base_url = base_url
         self.headers = {
             "Accept": "application/json",
             "User-agent": "leakix-client-python/%s" % __VERSION__,
         }
         if api_key:
             self.headers["api-key"] = api_key
 
-    def get(self, scope: Scope, queries: Optional[List[Query]] = None):
-        if queries is None or len(queries) == 0:
-            serialized_query = EmptyQuery().serialize()
-        else:
-            serialized_query = [q.serialize() for q in queries]
-            serialized_query = " ".join(serialized_query)
-            serialized_query = "%s" % serialized_query
-        url = "%s/search" % self.BASE_URL
+    def __get(self, url, params):
         r = requests.get(
             url,
-            params={"scope": scope.value, "q": serialized_query},
+            params=params,
             headers=self.headers,
         )
         if r.status_code == 200:
@@ -54,7 +51,18 @@ def get(self, scope: Scope, queries: Optional[List[Query]] = None):
         elif r.status_code == 429:
             return RateLimitResponse(response=r)
         else:
-            return ErrorResponse(response=r, response_json=response.json())
+            return ErrorResponse(response=r, response_json=r.json())
+
+    def get(self, scope: Scope, queries: Optional[List[Query]] = None):
+        if queries is None or len(queries) == 0:
+            serialized_query = EmptyQuery().serialize()
+        else:
+            serialized_query = [q.serialize() for q in queries]
+            serialized_query = " ".join(serialized_query)
+            serialized_query = "%s" % serialized_query
+        url = "%s/search" % self.base_url
+        r = self.__get(url=url, params={"scope": scope.value, "q": serialized_query})
+        return r
 
     def get_service(self, queries: Optional[List[Query]] = None):
         r = self.get(Scope.SERVICE, queries=queries)
@@ -73,18 +81,21 @@ def get_leak(self, queries: Optional[List[Query]] = None):
         return r
 
     def get_host(self, ipv4: str):
-        url = "%s/host/%s" % (self.BASE_URL, ipv4)
-        r = requests.get(url, headers=self.headers)
-        if r.status_code == 200:
+        url = "%s/host/%s" % (self.base_url, ipv4)
+        r = self.__get(url, params=None)
+        if r.is_success():
             response_json = r.json()
             formatted_result = HostResult.from_dict(response_json)
             response_json = {
                 "services": formatted_result.Services,
                 "leaks": formatted_result.Leaks,
             }
-            return SuccessResponse(response=r, response_json=response_json)
-        elif r.status_code == 429:
-            return RateLimitResponse(response=r)
-        else:
-            response_json = r.json()
-            return ErrorResponse(response=r, response_json=response_json)
+            r.response_json = response_json
+        return r
+
+    def get_plugins(self):
+        url = "%s/api/plugins" % (self.base_url)
+        r = self.__get(url, params=None)
+        if r.is_success():
+            r.response_json = [APIResult.from_dict(d) for d in r.json()]
+        return r

leakix/plugin.py

@@ -1,4 +1,10 @@
 from enum import Enum
+from serde import Model, fields
+
+
+class APIResult(Model):
+    name: fields.Str()
+    description: fields.Str()
 
 
 class Plugin(Enum):

leakix/response.py

@@ -4,14 +4,10 @@
 class AbstractResponse(metaclass=ABCMeta):
     def __init__(self, response, response_json=None):
         self.response = response
-        self.response_json = response_json
+        self.response_json = response_json or response.json()
 
     def json(self):
-        return (
-            self.response_json
-            if self.response_json is not None
-            else self.response.json()
-        )
+        return self.response_json
 
     def status_code(self):
         return self.response.status_code

tests/test_client.py

@@ -30,7 +30,7 @@ def test_get_host_success(client):
         # remove .json
         ipv4 = f.name[:-5]
         with requests_mock.Mocker() as m:
-            url = "%s/host/%s" % (client.BASE_URL, ipv4)
+            url = "%s/host/%s" % (client.base_url, ipv4)
             m.get(url, json=res_json, status_code=200)
             response = client.get_host(ipv4)
             assert response.is_success()
@@ -47,7 +47,7 @@ def test_get_host_404(client):
         ipv4 = f.name[:-5]
         client = Client()
         with requests_mock.Mocker() as m:
-            url = "%s/host/%s" % (client.BASE_URL, ipv4)
+            url = "%s/host/%s" % (client.base_url, ipv4)
             m.get(url, json=res_json, status_code=404)
             response = client.get_host(ipv4)
             assert response.is_error()
@@ -60,7 +60,7 @@ def test_get_host_429(client, fake_ipv4):
     status_code = 429
     res_json = {"reason": "rate-limit", "status": "error"}
     with requests_mock.Mocker() as m:
-        url = "%s/host/%s" % (client.BASE_URL, fake_ipv4)
+        url = "%s/host/%s" % (client.base_url, fake_ipv4)
         m.get(url, json=res_json, status_code=status_code)
         response = client.get_host(fake_ipv4)
         assert response.is_error()