diff --git a/backend/services/dns.py b/backend/services/dns.py index 937c2f7..b650da3 100644 --- a/backend/services/dns.py +++ b/backend/services/dns.py @@ -10,6 +10,8 @@ ares_query_cname_result, ares_query_txt_result, ) +from returns.future import future_safe +from returns.unsafe import unsafe_perform_io from backend import schemas from backend.schemas.dns import AAAA, CNAME, TXT, A @@ -20,17 +22,21 @@ QUERY_TYPES = typing.Literal["A", "AAAA", "CNAME", "TXT"] -async def query(name: str, query_type: QUERY_TYPES): - try: - resolver = aiodns.DNSResolver() - records = await resolver.query(name, query_type) +@future_safe +async def safe_query( + name: str, query_type: QUERY_TYPES, resolver: aiodns.DNSResolver | None = None +) -> list[typing.Any]: + resolver = resolver or aiodns.DNSResolver() + records = await resolver.query(name, query_type) - if not isinstance(records, list): - records = [records] + if not isinstance(records, list): + records = [records] - return typing.cast(list[typing.Any], records) - except aiodns.error.DNSError: - return [] + return typing.cast(list[typing.Any], records) + + +async def query(name: str, query_type: QUERY_TYPES) -> list[typing.Any]: + return unsafe_perform_io((await safe_query(name, query_type)).value_or([])) async def query_a_records(name: str) -> list[A]: diff --git a/backend/services/fingerprint.py b/backend/services/fingerprint.py index 47d83cf..e7dd538 100644 --- a/backend/services/fingerprint.py +++ b/backend/services/fingerprint.py @@ -1,6 +1,6 @@ -import contextlib -from dataclasses import dataclass, field +from dataclasses import dataclass +import aiometer import httpx from returns.functions import raise_exception from returns.future import FutureResultE, future_safe @@ -24,102 +24,96 @@ class Container: response: httpx.Response - html: schemas.HTML | None = field(default=None) - dns: schemas.DNS | None = field(default=None) - tracker: schemas.Tracker | None = field(default=None) - whois: schemas.Whois | None = field(default=None) - favicon: schemas.Favicon | None = field(default=None) - certificate: schemas.Certificate | None = field(default=None) - tls: schemas.TLS | None = field(default=None) - @future_safe -async def get_url(url: str) -> Container: +async def get_url(url: str) -> httpx.Response: async with httpx.AsyncClient(verify=False) as client: - response = await client.get(url, follow_redirects=True) - return Container(response=response) + return await client.get(url, follow_redirects=True) @future_safe -async def get_whois(container: Container) -> Container: - container.whois = await Whois().call(container.response) - return container +async def get_whois(res: httpx.Response) -> schemas.Whois: + return await Whois().call(res) @future_safe -async def get_certificate(container: Container) -> Container: - with contextlib.suppress(Exception): - container.certificate = await Certificate().call(container.response) - - return container +async def get_certificate(res: httpx.Response) -> schemas.Certificate | None: + return await Certificate().call(res) @future_safe -async def get_favicon(container: Container) -> Container: - with contextlib.suppress(Exception): - container.favicon = await Favicon().call(container.response) - - return container +async def get_favicon(res: httpx.Response) -> schemas.Favicon | None: + return await Favicon().call(res) @future_safe -async def get_html(container: Container) -> Container: - with contextlib.suppress(Exception): - container.html = await HTML().call(container.response) - - return container +async def get_html(res: httpx.Response) -> schemas.HTML: + return await HTML().call(res) @future_safe -async def get_tracker(container: Container) -> Container: - try: - container.tracker = await Tracker().call(container.response) - except Exception: - container.tracker = schemas.Tracker() - - return container +async def get_tracker(res: httpx.Response) -> schemas.Tracker: + return await Tracker().call(res) @future_safe -async def get_dns(container: Container) -> Container: - try: - container.dns = await DNS().call(container.response) - except Exception: - container.dns = schemas.DNS() - - return container +async def get_dns(res: httpx.Response) -> schemas.DNS: + return await DNS().call(res) @future_safe -async def get_tls(container: Container) -> Container: - with contextlib.suppress(Exception): - container.tls = await TLS().call(container.response) +async def get_tls(res: httpx.Response) -> schemas.TLS | None: + return await TLS().call(res) + - return container +@future_safe +async def get_fingerprint(res: httpx.Response) -> schemas.Fingerprint: + ( + certificate_result, + dns_result, + favicon_result, + html_result, + tls_result, + tracker_result, + whois_result, + ) = await aiometer.run_all( + [ + get_certificate(res).awaitable, + get_dns(res).awaitable, + get_favicon(res).awaitable, + get_html(res).awaitable, + get_tls(res).awaitable, + get_tracker(res).awaitable, + get_whois(res).awaitable, + ] + ) + certificate: schemas.Certificate | None = unsafe_perform_io( + certificate_result.value_or(None) # type: ignore + ) # type: ignore + dns: schemas.DNS = unsafe_perform_io(dns_result.value_or(schemas.DNS())) # type: ignore + favicon: schemas.Favicon = unsafe_perform_io(favicon_result.value_or(None)) # type: ignore + html: schemas.HTML = unsafe_perform_io(html_result.value_or(None)) # type: ignore + tls: schemas.TLS = unsafe_perform_io(tls_result.value_or(None)) # type: ignore + tracker: schemas.Tracker = unsafe_perform_io( + tracker_result.value_or(schemas.Tracker()) # type: ignore + ) + whois: schemas.Whois = unsafe_perform_io(whois_result.value_or(schemas.Whois())) # type: ignore + return schemas.Fingerprint( + certificate=certificate, + dns=dns, + favicon=favicon, + html=html, + tls=tls, + tracker=tracker, + whois=whois, + headers=dict(res.headers), + ) class Fingerprint(AbstractService): async def call(self, url: str) -> schemas.Fingerprint: - f_result: FutureResultE[Container] = flow( - url, - get_url, - bind(get_certificate), - bind(get_dns), - bind(get_favicon), - bind(get_html), - bind(get_tls), - bind(get_tracker), - bind(get_whois), + f_result: FutureResultE[schemas.Fingerprint] = flow( + url, get_url, bind(get_fingerprint) ) result = (await f_result.awaitable()).alt(raise_exception).unwrap() - container = unsafe_perform_io(result) - return schemas.Fingerprint( - html=container.html, # type: ignore - dns=container.dns, # type: ignore - tracker=container.tracker, # type: ignore - whois=container.whois, # type: ignore - favicon=container.favicon, - certificate=container.certificate, - tls=container.tls, - headers=dict(container.response.headers), - ) + return unsafe_perform_io(result)