From 28056f6a9175b8a041123fa68f2a338c6093075e Mon Sep 17 00:00:00 2001 From: Manabu Niseki Date: Sat, 9 Mar 2024 12:13:07 +0900 Subject: [PATCH 1/2] refactor: make EmalRep optional --- backend/api/endpoints/analyze.py | 34 ++++++++++++++-------------- backend/api/endpoints/lookup.py | 6 +++-- backend/api/endpoints/submit.py | 6 ++--- backend/clients/emailrep.py | 8 +++++-- backend/{deps.py => dependencies.py} | 15 +++++++----- backend/factories/response.py | 10 ++++---- backend/settings.py | 3 +++ 7 files changed, 47 insertions(+), 35 deletions(-) rename backend/{deps.py => dependencies.py} (85%) diff --git a/backend/api/endpoints/analyze.py b/backend/api/endpoints/analyze.py index 9a6e239..b1129af 100644 --- a/backend/api/endpoints/analyze.py +++ b/backend/api/endpoints/analyze.py @@ -3,7 +3,7 @@ from pydantic import ValidationError from redis import Redis -from backend import clients, deps, schemas, settings +from backend import clients, dependencies, schemas, settings from backend.factories.response import ResponseFactory router = APIRouter() @@ -13,7 +13,7 @@ async def _analyze( file: bytes, *, spam_assassin: clients.SpamAssassin, - email_rep: clients.EmailRep, + optional_email_rep: clients.EmailRep | None = None, optional_inquest: clients.InQuest | None = None, optional_vt: clients.VirusTotal | None = None, optional_urlscan: clients.UrlScan | None = None, @@ -28,7 +28,7 @@ async def _analyze( return await ResponseFactory.call( payload.file, - email_rep=email_rep, + optional_email_rep=optional_email_rep, spam_assassin=spam_assassin, optional_inquest=optional_inquest, optional_urlscan=optional_urlscan, @@ -56,17 +56,17 @@ async def analyze( payload: schemas.Payload, *, 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, + spam_assassin: dependencies.SpamAssassin, + optional_redis: dependencies.OptionalRedis, + optional_email_rep: dependencies.OptionalEmailRep, + optional_inquest: dependencies.OptionalInQuest, + optional_vt: dependencies.OptionalVirusTotal, + optional_urlscan: dependencies.OptionalUrlScan, ) -> schemas.Response: response = await _analyze( payload.file.encode(), - email_rep=email_rep, spam_assassin=spam_assassin, + optional_email_rep=optional_email_rep, optional_inquest=optional_inquest, optional_urlscan=optional_urlscan, optional_vt=optional_vt, @@ -90,16 +90,16 @@ async def analyze_file( file: bytes = 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, + optional_redis: dependencies.OptionalRedis, + spam_assassin: dependencies.SpamAssassin, + optional_email_rep: dependencies.OptionalEmailRep, + optional_inquest: dependencies.OptionalInQuest, + optional_vt: dependencies.OptionalVirusTotal, + optional_urlscan: dependencies.OptionalUrlScan, ) -> schemas.Response: response = await _analyze( file, - email_rep=email_rep, + optional_email_rep=optional_email_rep, spam_assassin=spam_assassin, optional_inquest=optional_inquest, optional_urlscan=optional_urlscan, diff --git a/backend/api/endpoints/lookup.py b/backend/api/endpoints/lookup.py index e9c539b..c4eea01 100644 --- a/backend/api/endpoints/lookup.py +++ b/backend/api/endpoints/lookup.py @@ -1,6 +1,6 @@ from fastapi import APIRouter, HTTPException, status -from backend import deps, schemas, settings +from backend import dependencies, schemas, settings router = APIRouter() @@ -11,7 +11,9 @@ summary="Lookup cached analysis", description="Try to fetch existing analysis from database", ) -async def lookup(id: str, *, optional_redis: deps.OptionalRedis) -> schemas.Response: +async def lookup( + id: str, *, optional_redis: dependencies.OptionalRedis +) -> schemas.Response: if optional_redis is None: raise HTTPException( status_code=status.HTTP_501_NOT_IMPLEMENTED, diff --git a/backend/api/endpoints/submit.py b/backend/api/endpoints/submit.py index 82b27ff..22cd890 100644 --- a/backend/api/endpoints/submit.py +++ b/backend/api/endpoints/submit.py @@ -1,7 +1,7 @@ import httpx from fastapi import APIRouter, HTTPException, status -from backend import deps, schemas +from backend import dependencies, schemas from backend.schemas.eml import Attachment from backend.utils import attachment_to_file @@ -16,7 +16,7 @@ status_code=200, ) async def submit_to_inquest( - attachment: Attachment, *, optional_inquest: deps.OptionalInQuest + attachment: Attachment, *, optional_inquest: dependencies.OptionalInQuest ) -> schemas.SubmissionResult: # check ext type valid_types = ["doc", "docx", "ppt", "pptx", "xls", "xlsx"] @@ -49,7 +49,7 @@ async def submit_to_inquest( status_code=200, ) async def submit_to_virustotal( - attachment: Attachment, *, optional_vt: deps.OptionalVirusTotal + attachment: Attachment, *, optional_vt: dependencies.OptionalVirusTotal ) -> schemas.SubmissionResult: if optional_vt is None: raise HTTPException( diff --git a/backend/clients/emailrep.py b/backend/clients/emailrep.py index 6a568f3..3580d01 100644 --- a/backend/clients/emailrep.py +++ b/backend/clients/emailrep.py @@ -1,11 +1,15 @@ import httpx +from starlette.datastructures import Secret from backend import schemas class EmailRep(httpx.AsyncClient): - def __init__(self) -> None: - super().__init__(base_url="https://emailrep.io") + def __init__(self, api_key: Secret) -> None: + super().__init__( + base_url="https://emailrep.io", + headers={"key": str(api_key), "user-agent": "EML-Analyzer"}, + ) async def lookup(self, email: str) -> schemas.EmailRepLookup: r = await self.get(f"/{email}") diff --git a/backend/deps.py b/backend/dependencies.py similarity index 85% rename from backend/deps.py rename to backend/dependencies.py index 785e853..ff323f5 100644 --- a/backend/deps.py +++ b/backend/dependencies.py @@ -71,13 +71,16 @@ async def get_optional_urlscan(): @asynccontextmanager -async def _get_email_rep(): - async with clients.EmailRep() as client: - yield client +async def _get_optional_email_rep(api_key: Secret | None = settings.EMAIL_REP_API_KEY): + if api_key is None: + yield None + else: + async with clients.EmailRep(api_key=api_key) as client: + yield client -async def get_email_rep(): - async with _get_email_rep() as client: +async def get_optional_email_rep(): + async with _get_optional_email_rep(settings.EMAIL_REP_API_KEY) as client: yield client @@ -101,5 +104,5 @@ def get_spam_assassin() -> clients.SpamAssassin: clients.UrlScan | None, Depends(get_optional_urlscan) ] -EmailRep = typing.Annotated[clients.EmailRep, Depends(get_email_rep)] +OptionalEmailRep = typing.Annotated[clients.EmailRep, Depends(get_optional_email_rep)] SpamAssassin = typing.Annotated[clients.SpamAssassin, Depends(get_spam_assassin)] diff --git a/backend/factories/response.py b/backend/factories/response.py index d58fe84..5b8faa7 100644 --- a/backend/factories/response.py +++ b/backend/factories/response.py @@ -75,8 +75,8 @@ async def set_verdicts( response: schemas.Response, *, eml_file: bytes, - email_rep: clients.EmailRep, spam_assassin: clients.SpamAssassin, + optional_email_rep: clients.EmailRep | None = None, optional_vt: clients.VirusTotal | None = None, optional_urlscan: clients.UrlScan | None = None, optional_inquest: clients.InQuest | None = None, @@ -86,9 +86,9 @@ async def set_verdicts( get_oleid_verdict(response.eml.attachments), ] - if response.eml.header.from_ is not None: + if response.eml.header.from_ is not None and optional_email_rep is not None: f_results.append( - get_email_rep_verdicts(response.eml.header.from_, client=email_rep) + get_email_rep_verdicts(response.eml.header.from_, client=optional_email_rep) ) if optional_vt is not None: @@ -117,8 +117,8 @@ async def call( cls, eml_file: bytes, *, - email_rep: clients.EmailRep, spam_assassin: clients.SpamAssassin, + optional_email_rep: clients.EmailRep | None, optional_vt: clients.VirusTotal | None = None, optional_urlscan: clients.UrlScan | None = None, optional_inquest: clients.InQuest | None = None, @@ -129,7 +129,7 @@ async def call( partial( set_verdicts, eml_file=eml_file, - email_rep=email_rep, + optional_email_rep=optional_email_rep, spam_assassin=spam_assassin, optional_vt=optional_vt, optional_urlscan=optional_urlscan, diff --git a/backend/settings.py b/backend/settings.py index dc3bab3..3c2aeab 100644 --- a/backend/settings.py +++ b/backend/settings.py @@ -35,6 +35,9 @@ ) INQUEST_API_KEY: Secret | None = config("INQUEST_API_KEY", cast=Secret, default=None) URLSCAN_API_KEY: Secret | None = config("URLSCAN_API_KEY", cast=Secret, default=None) +EMAIL_REP_API_KEY: Secret | None = config( + "EMAIL_REP_API_KEY", cast=Secret, default=None +) # Async/aiometer ASYNC_MAX_AT_ONCE: int | None = config("ASYNC_MAX_AT_ONCE", cast=int, default=None) From 1db97d23354f6e07d56c422ddae282245d009b57 Mon Sep 17 00:00:00 2001 From: Manabu Niseki Date: Sat, 9 Mar 2024 13:22:44 +0900 Subject: [PATCH 2/2] fix: rename module --- backend/api/endpoints/cache.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backend/api/endpoints/cache.py b/backend/api/endpoints/cache.py index c1feb40..afd18d0 100644 --- a/backend/api/endpoints/cache.py +++ b/backend/api/endpoints/cache.py @@ -1,6 +1,6 @@ from fastapi import APIRouter, HTTPException, status -from backend import deps, settings +from backend import dependencies, settings router = APIRouter() @@ -11,7 +11,7 @@ summary="Get analysis cache keys", description="Try to get analysis cache keys", ) -async def cache_keys(optional_redis: deps.OptionalRedis) -> list[str]: +async def cache_keys(optional_redis: dependencies.OptionalRedis) -> list[str]: if optional_redis is None: raise HTTPException( status_code=status.HTTP_501_NOT_IMPLEMENTED,