Merge pull request #204 from ninoseki/introduce-returns

refactor: introduce returns

.github/workflows/python.yml

@@ -45,6 +45,11 @@ jobs:
         run: poetry run pyupgrade --py311-plus **/*.py
   test:
     runs-on: ubuntu-latest
+    services:
+      spamassassin:
+        image: instantlinux/spamassassin:4.0.0-6
+        ports:
+          - 783:783
     strategy:
       matrix:
         python-version: [3.11]
@@ -63,6 +68,8 @@ jobs:
         run: poetry install
       - name: Make dummy frontend directory
         run: mkdir -p frontend/dist/
+      - name: Wait until SpamAssasin ready
+        run: poetry run python scripts/ping.py
       - name: Run tests
         run: poetry run pytest -v --cov=app --cov-report=term-missing
       - name: Coveralls

backend/__init__.py

@@ -0,0 +1,3 @@
+from .clients.emailrep import EmailRep  # noqa: F401
+from .clients.inquest import InQuest  # noqa: F401
+from .clients.spamassasin import SpamAssassin  # noqa: F401

backend/api/endpoints/analyze.py

@@ -3,14 +3,21 @@
 from pydantic import ValidationError
 from redis import Redis
 
-from backend import deps, schemas
-from backend.core import settings
+from backend import clients, deps, schemas, settings
 from backend.factories.response import ResponseFactory
 
 router = APIRouter()
 
 
-async def _analyze(file: bytes) -> schemas.Response:
+async def _analyze(
+    file: bytes,
+    *,
+    spam_assassin: clients.SpamAssassin,
+    email_rep: clients.EmailRep,
+    optional_inquest: clients.InQuest | None = None,
+    optional_vt: clients.VirusTotal | None = None,
+    optional_urlscan: clients.UrlScan | None = None,
+) -> schemas.Response:
     try:
         payload = schemas.FilePayload(file=file)
     except ValidationError as exc:
@@ -19,7 +26,14 @@ async def _analyze(file: bytes) -> schemas.Response:
             detail=jsonable_encoder(exc.errors()),
         ) from exc
 
-    return await ResponseFactory.from_bytes(payload.file)
+    return await ResponseFactory.call(
+        payload.file,
+        email_rep=email_rep,
+        spam_assassin=spam_assassin,
+        optional_inquest=optional_inquest,
+        optional_urlscan=optional_urlscan,
+        optional_vt=optional_vt,
+    )
 
 
 def cache_response(
@@ -45,8 +59,20 @@ async def analyze(
     *,
     background_tasks: BackgroundTasks,
     optional_redis: deps.OptionalRedis,
+    spam_assassin: deps.SpamAssassin,
+    email_rep: deps.EmailRep,
+    optional_inquest: deps.OptionalInQuest,
+    optional_vt: deps.OptionalVirusTotal,
+    optional_urlscan: deps.OptionalUrlScan,
 ) -> schemas.Response:
-    response = await _analyze(payload.file.encode())
+    response = await _analyze(
+        payload.file.encode(),
+        email_rep=email_rep,
+        spam_assassin=spam_assassin,
+        optional_inquest=optional_inquest,
+        optional_urlscan=optional_urlscan,
+        optional_vt=optional_vt,
+    )
 
     if optional_redis is not None:
         background_tasks.add_task(
@@ -67,8 +93,20 @@ async def analyze_file(
     *,
     background_tasks: BackgroundTasks,
     optional_redis: deps.OptionalRedis,
+    spam_assassin: deps.SpamAssassin,
+    email_rep: deps.EmailRep,
+    optional_inquest: deps.OptionalInQuest,
+    optional_vt: deps.OptionalVirusTotal,
+    optional_urlscan: deps.OptionalUrlScan,
 ) -> schemas.Response:
-    response = await _analyze(file)
+    response = await _analyze(
+        file,
+        email_rep=email_rep,
+        spam_assassin=spam_assassin,
+        optional_inquest=optional_inquest,
+        optional_urlscan=optional_urlscan,
+        optional_vt=optional_vt,
+    )
 
     if optional_redis is not None:
         background_tasks.add_task(

backend/api/endpoints/lookup.py

@@ -1,7 +1,6 @@
 from fastapi import APIRouter, HTTPException, status
 
-from backend import deps, schemas
-from backend.core import settings
+from backend import deps, schemas, settings
 
 router = APIRouter()
 

backend/api/endpoints/submit.py

@@ -1,11 +1,9 @@
+import httpx
 from fastapi import APIRouter, HTTPException, status
-from httpx._exceptions import HTTPError
 
-from backend.core.utils import has_inquest_api_key, has_virustotal_api_key
+from backend import deps, schemas
 from backend.schemas.eml import Attachment
-from backend.schemas.submission import SubmissionResult
-from backend.submitters.inquest import InQuestSubmitter
-from backend.submitters.virustotal import VirusTotalSubmitter
+from backend.utils import attachment_to_file
 
 router = APIRouter()
 
@@ -17,12 +15,9 @@
     description="Submit an attachment to InQuest",
     status_code=200,
 )
-async def submit_to_inquest(attachment: Attachment) -> SubmissionResult:
-    if not has_inquest_api_key():
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="You don't have the InQuest API key",
-        )
+async def submit_to_inquest(
+    attachment: Attachment, *, optional_inquest: deps.OptionalInQuest
+) -> schemas.SubmissionResult:
     # check ext type
     valid_types = ["doc", "docx", "ppt", "pptx", "xls", "xlsx"]
     if attachment.extension not in valid_types:
@@ -31,36 +26,45 @@ async def submit_to_inquest(attachment: Attachment) -> SubmissionResult:
             detail=f"{attachment.extension} is not supported.",
         )
 
-    submitter = InQuestSubmitter(attachment)
+    if optional_inquest is None:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="You don't have the InQuest API key",
+        )
+
     try:
-        return await submitter.submit()
-    except HTTPError as e:
+        return await optional_inquest.submit(attachment_to_file(attachment))
+    except httpx.HTTPStatusError as e:
         raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Something went wrong with InQuest submission: {e!s}",
+            status_code=e.response.status_code,
+            detail=f"Something went wrong with InQuest submission: {e}",
         ) from e
 
 
 @router.post(
     "/virustotal",
-    response_model=SubmissionResult,
     response_description="Return a submission result",
     summary="Submit an attachment to VirusTotal",
     description="Submit an attachment to VirusTotal",
     status_code=200,
 )
-async def submit_to_virustotal(attachment: Attachment) -> SubmissionResult:
-    if not has_virustotal_api_key():
+async def submit_to_virustotal(
+    attachment: Attachment, *, optional_vt: deps.OptionalVirusTotal
+) -> schemas.SubmissionResult:
+    if optional_vt is None:
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail="You don't have the VirusTotal API key",
         )
 
-    submitter = VirusTotalSubmitter(attachment)
     try:
-        return await submitter.submit()
-    except HTTPError as e:
+        await optional_vt.scan_file_async(attachment_to_file(attachment))
+        sha256 = attachment.hash.sha256
+        return schemas.SubmissionResult(
+            reference_url=f"https://www.virustotal.com/gui/file/{sha256}/detection"
+        )
+    except httpx.HTTPStatusError as e:
         raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Something went wrong with VirusTotal submission: {e!s}",
+            status_code=e.response.status_code,
+            detail=f"Something went wrong with VirusTotal submission: {e}",
         ) from e

backend/clients/__init__.py

@@ -0,0 +1,8 @@
+import vt
+
+from .emailrep import EmailRep  # noqa: F401
+from .inquest import InQuest  # noqa: F401
+from .spamassasin import SpamAssassin  # noqa: F401
+from .urlscan import UrlScan  # noqa: F401
+
+VirusTotal = vt.Client

backend/clients/emailrep.py

@@ -0,0 +1,13 @@
+import httpx
+
+from backend import schemas
+
+
+class EmailRep(httpx.AsyncClient):
+    def __init__(self) -> None:
+        super().__init__(base_url="https://emailrep.io")
+
+    async def lookup(self, email: str) -> schemas.EmailRepLookup:
+        r = await self.get(f"/{email}")
+        r.raise_for_status()
+        return schemas.EmailRepLookup.model_validate(r.json())

backend/clients/inquest.py

@@ -0,0 +1,27 @@
+import io
+
+import httpx
+from starlette.datastructures import Secret
+
+from backend import schemas
+
+
+class InQuest(httpx.AsyncClient):
+    def __init__(self, api_key: Secret) -> None:
+        super().__init__(
+            base_url="https://labs.inquest.net",
+            headers={"Authorization": f"Basic: {api_key}"},
+        )
+
+    async def lookup(self, sha256: str) -> schemas.InQuestLookup:
+        r = await self.get("/api/dfi/details", params={"sha256": sha256})
+        r.raise_for_status()
+        return schemas.InQuestLookup.model_validate(r.json())
+
+    async def submit(self, f: io.BytesIO) -> schemas.SubmissionResult:
+        r = await self.post("/api/dfi/upload", files={"file": f})
+        r.raise_for_status()
+        data = r.json()["data"]
+        return schemas.SubmissionResult(
+            reference_url=f"https://labs.inquest.net/dfi/sha256/{data}"
+        )

backend/services/spamassassin.py → backend/clients/spamassasin.py

@@ -1,46 +1,32 @@
-from dataclasses import dataclass
-from typing import Any
-
 import aiospamc
+from aiospamc.header_values import Headers
 from async_timeout import timeout
 
-
-@dataclass
-class Detail:
-    name: str
-    score: float
-    description: str
-
-
-@dataclass
-class Report:
-    score: float
-    details: list[Detail]
-    level = 5.0
-
-    def is_spam(self, level: float = 5.0) -> bool:
-        return self.score is None or self.score > level
+from backend import schemas, settings
 
 
 class Parser:
-    def __init__(self, headers: dict[str, Any], body: str):
+    def __init__(self, headers: Headers, body: str):
         self.headers = headers
         self.body = body
         self.score = 0.0
-        self.details: list[Detail] = []
+        self.details: list[schemas.SpamAssassinDetail] = []
 
     def _parse_headers(self):
-        spam_value = self.headers["Spam"]
-        self.score = spam_value.score
+        spam_value = self.headers.get("Spam")
+        if spam_value is not None:
+            self.score = spam_value.score
 
-    def _parse_detail(self, line: str) -> Detail:
+    def _parse_detail(self, line: str) -> schemas.SpamAssassinDetail:
         parts = line.split()
         score = float(parts[0])
         name = parts[1]
         description = " ".join(parts[2:])
-        return Detail(name=name, score=score, description=description)
+        return schemas.SpamAssassinDetail(
+            name=name, score=score, description=description
+        )
 
-    def _parse_details(self, details: str) -> list[Detail]:
+    def _parse_details(self, details: str) -> list[schemas.SpamAssassinDetail]:
         lines = details.splitlines()
         normalized_line: list[str] = []
 
@@ -56,33 +42,35 @@ def _parse_details(self, details: str) -> list[Detail]:
 
     def _parse_body(self):
         lines = [line for line in self.body.splitlines() if line != ""]
-        demiliter_index = 0
+        delimiter_index = 0
         for index, line in enumerate(lines):
             if "---" in line:
-                demiliter_index = index + 1
+                delimiter_index = index + 1
                 break
 
-        details = "\n".join(lines[demiliter_index:])
+        details = "\n".join(lines[delimiter_index:])
         self.details = self._parse_details(details)
 
-    def parse(self):
+    def parse(self) -> schemas.SpamAssassinReport:
         self._parse_headers()
         self._parse_body()
-
-    def to_report(self) -> Report:
-        return Report(score=self.score, details=self.details)
+        return schemas.SpamAssassinReport(score=self.score, details=self.details)
 
 
 class SpamAssassin:
-    def __init__(self, host: str = "127.0.0.1", port: int = 783, timeout: int = 10):
+    def __init__(
+        self,
+        host: str = settings.SPAMASSASSIN_HOST,
+        port: int = settings.SPAMASSASSIN_PORT,
+        timeout: int = settings.SPAMASSASSIN_TIMEOUT,
+    ):
         self.host = host
         self.port = port
         self.timeout = timeout
 
-    async def report(self, message: bytes) -> Report:
+    async def report(self, message: bytes) -> schemas.SpamAssassinReport:
         async with timeout(self.timeout):
             response = await aiospamc.report(message, host=self.host, port=self.port)
 
         parser = Parser(headers=response.headers, body=response.body.decode())
-        parser.parse()
-        return parser.to_report()
+        return parser.parse()

backend/clients/urlscan.py

@@ -0,0 +1,26 @@
+from urllib.parse import urlparse
+
+import httpx
+from starlette.datastructures import Secret
+
+from backend import schemas
+
+
+class UrlScan(httpx.AsyncClient):
+    def __init__(self, api_key: Secret) -> None:
+        super().__init__(
+            base_url="https://urlscan.io", headers={"api-key": str(api_key)}
+        )
+
+    async def lookup(
+        self,
+        url: str,
+    ) -> schemas.UrlScanLookup:
+        parsed = urlparse(url)
+        params = {
+            "q": f'task.url:"{url}" AND task.domain:"{parsed.hostname}" AND verdicts.malicious:true',
+            "size": 1,
+        }
+        r = await self.get("/api/v1/search/", params=params)
+        r.raise_for_status()
+        return schemas.UrlScanLookup.model_validate(r.json())