Merge pull request #214 from ninoseki/query-lookup

feat: quick query lookup

backend/api/api.py

@@ -1,8 +1,9 @@
 from fastapi import APIRouter
 
-from backend.api.endpoints import fingerprint
+from backend.api.endpoints import counters, fingerprint
 
 api_router = APIRouter()
 api_router.include_router(
     fingerprint.router, prefix="/fingerprint", tags=["fingerprint"]
 )
+api_router.include_router(counters.router, prefix="/counters", tags=["counters"])

backend/api/endpoints/counters.py

@@ -0,0 +1,41 @@
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, HTTPException, status
+
+from backend import schemas, services, settings
+
+router = APIRouter()
+
+
+def get_shodan():
+    if settings.SHODAN_API_KEY is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Shodan is not configured correctly",
+        )
+
+    return services.Shodan(settings.SHODAN_API_KEY)
+
+
+def get_censys():
+    if settings.CENSYS_APP_ID is None or settings.CENSYS_SECRET is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Censys is not configured correctly",
+        )
+
+    return services.Censys(settings.CENSYS_APP_ID, settings.CENSYS_SECRET)
+
+
+Shodan = Annotated[services.Shodan, Depends(get_shodan)]
+Censys = Annotated[services.Censys, Depends(get_censys)]
+
+
+@router.get("/shodan")
+async def shodan(query: str, *, shodan: Shodan) -> schemas.Count:
+    return schemas.Count(count=await shodan.call(query))
+
+
+@router.get("/censys")
+async def censys(query: str, *, censys: Censys):
+    return schemas.Count(count=await censys.call(query))

backend/main.py

@@ -3,8 +3,8 @@
 from fastapi.staticfiles import StaticFiles
 from loguru import logger
 
+from backend import settings
 from backend.api.api import api_router
-from backend.core import settings
 
 
 def create_app():

backend/schemas/__init__.py

@@ -1,3 +1,5 @@
+from pydantic import BaseModel
+
 from .certificate import Certificate  # noqa: F401
 from .dns import DNS  # noqa: F401
 from .favicon import Favicon  # noqa: F401
@@ -7,3 +9,7 @@
 from .tls import TLS  # noqa: F401
 from .tracker import Tracker  # noqa: F401
 from .whois import Whois  # noqa: F401
+
+
+class Count(BaseModel):
+    count: int

backend/services/__init__.py

@@ -1,4 +1,5 @@
 from .certificate import Certificate  # noqa: F401
+from .counters import Censys, Shodan  # noqa: F401
 from .dns import DNS  # noqa: F401
 from .favicon import Favicon  # noqa: F401
 from .fingerprint import Fingerprint  # noqa: F401

backend/services/counters.py

@@ -0,0 +1,37 @@
+import httpx
+from starlette.datastructures import Secret
+
+
+class Shodan:
+    def __init__(self, api_key: Secret):
+        self.api_key = api_key
+
+    async def call(self, query: str):
+        async with httpx.AsyncClient(base_url="https://api.shodan.io") as client:
+            res = await client.get(
+                "/shodan/host/search", params={"key": str(self.api_key), "query": query}
+            )
+            res.raise_for_status()
+
+        return res.json()["total"]
+
+
+class Censys:
+    def __init__(
+        self,
+        app_id: Secret,
+        secret: Secret,
+    ):
+        self.app_id = app_id
+        self.secret = secret
+
+    async def call(self, query: str):
+        auth = httpx.BasicAuth(username=str(self.app_id), password=str(self.secret))
+
+        async with httpx.AsyncClient(
+            base_url="https://search.censys.io", auth=auth
+        ) as client:
+            res = await client.get("/api/v2/hosts/search", params={"q": query})
+            res.raise_for_status()
+
+        return res.json()["result"]["total"]

backend/core/settings.py → backend/settings.py

@@ -1,11 +1,9 @@
 import sys
 
 from starlette.config import Config
+from starlette.datastructures import Secret
 
-try:
-    config = Config(".env")
-except Exception:
-    config = Config()
+config = Config(".env")
 
 PROJECT_NAME: str = config("PROJECT_NAME", default="shimon")
 
@@ -15,3 +13,8 @@
 LOG_FILE = config("LOG_FILE", default=sys.stderr)
 LOG_LEVEL: str = config("LOG_LEVEL", cast=str, default="DEBUG")
 LOG_BACKTRACE: bool = config("LOG_BACKTRACE", cast=bool, default=True)
+
+
+SHODAN_API_KEY: Secret | None = config("SHODAN_API_KEY", cast=Secret, default=None)
+CENSYS_APP_ID: Secret | None = config("CENSYS_APP_ID", cast=Secret, default=None)
+CENSYS_SECRET: Secret | None = config("CENSYS_SECRET", cast=Secret, default=None)