diff --git a/pyproject.toml b/pyproject.toml index b00e15d..14efa9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ license = "Apache-2.0" license-files = ["LICENSE"] readme = "README.md" requires-python = ">=3.11" -version = "2.1.1" +version = "2.2.0" [project.optional-dependencies] cli = [ diff --git a/src/mega/api.py b/src/mega/api.py index 07cc3df..6a12a46 100644 --- a/src/mega/api.py +++ b/src/mega/api.py @@ -3,17 +3,19 @@ import asyncio import contextlib import dataclasses +import json import logging +import uuid from collections.abc import Mapping, Sequence from contextvars import ContextVar from functools import wraps -from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Literal, ParamSpec, Self, TypeVar import aiohttp import yarl from aiolimiter import AsyncLimiter +from mega import __version__, _package_name_ from mega.crypto import generate_hashcash from mega.errors import RequestError, RetryRequestError from mega.utils import random_id, random_u32int @@ -25,11 +27,8 @@ _R = TypeVar("_R") -_UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:137.0) Gecko/20100101 Firefox/137.0" -_DEFAULT_HEADERS: MappingProxyType[str, str] = MappingProxyType({"User-Agent": _UA}) - - LOG_HTTP_TRAFFIC: ContextVar[bool] = ContextVar("LOG_HTTP_TRAFFIC", default=False) + logger = logging.getLogger(__name__) @@ -78,14 +77,16 @@ class MegaAPI: _rate_limiter: AsyncLimiter _entrypoint: ClassVar[yarl.URL] = yarl.URL("https://g.api.mega.co.nz/cs") + user_agent: str - def __init__(self, session: aiohttp.ClientSession | None = None) -> None: + def __init__(self, session: aiohttp.ClientSession | None = None, user_agent: str | None = None) -> None: self.session_id = None self._request_id = random_u32int() self._client_id = random_id(10) self.__session = session self._auto_close_session = session is None self._rate_limiter = AsyncLimiter(100, 60) + self.user_agent = user_agent or f"{_package_name_}/{__version__}" @property def entrypoint(self) -> yarl.URL: @@ -174,19 +175,37 @@ async def __request( headers: Mapping[str, str] | None = None, **kwargs: Any, ) -> AsyncGenerator[aiohttp.ClientResponse]: - kwargs["headers"] = _DEFAULT_HEADERS | (headers or {}) + kwargs["headers"] = {"User-Agent": self.user_agent} | (headers or {}) + request_id = str(uuid.uuid4()) if LOG_HTTP_TRAFFIC.get(): - params = ", ".join(f"{name} = {value!r}" for name, value in kwargs.items()) - logger.debug(f"Making {method} request to {url!s} with {params}") - async with self._rate_limiter, self._session.request(method, url, **kwargs) as resp: - yield resp + logger.debug( + "Starting %s request [id=%s] to %s \n%s", + method, + request_id, + url, + kwargs, + ) + + resp = None + try: + async with self._rate_limiter, self._session.request(method, url, **kwargs) as resp: + yield resp + except RetryRequestError: + logger.warning("Request [id=%s] failed, retrying", request_id) + raise + finally: + if resp and LOG_HTTP_TRAFFIC.get(): + logger.debug( + "Finished %s request [id=%s]\n%s", + method, + request_id, + _LazyResponseLog(resp), + ) @staticmethod async def _parse_response(response: aiohttp.ClientResponse) -> Any: json_resp = await response.json() resp = json_resp - if LOG_HTTP_TRAFFIC.get(): - logger.debug(f"Got response [{response.status}] json={json_resp!r}") if isinstance(json_resp, list) and len(json_resp) == 1: resp = json_resp[0] @@ -196,19 +215,40 @@ async def _parse_response(response: aiohttp.ClientResponse) -> Any: return resp if resp == -3: - msg = "Request failed, retrying" - logger.warning(msg) raise RetryRequestError raise RequestError(resp) return resp +class _LazyResponseLog: + def __init__(self, resp: aiohttp.ClientResponse) -> None: + self.resp = resp + + def __json__(self) -> dict[str, Any]: + me = { + "url": str(self.resp.url), + "status_code": self.resp.status, + "response_headers": dict(self.resp.headers), + "content": None, + } + if self.resp._body: + stripped = self.resp._body.strip() + if stripped: + content = json.loads(stripped.decode(self.resp.get_encoding())) + me.update(content=content) + + return me + + def __str__(self) -> str: + return str(self.__json__()) + + class APIContextManager: __slots__ = ("_api",) - def __init__(self, session: aiohttp.ClientSession | None = None) -> None: - self._api: MegaAPI = MegaAPI(session) + def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None: + self._api: MegaAPI = MegaAPI(session, user_agent) async def __aenter__(self) -> Self: return self diff --git a/src/mega/client.py b/src/mega/client.py index d997074..2c2dfbb 100644 --- a/src/mega/client.py +++ b/src/mega/client.py @@ -43,8 +43,8 @@ class MegaNzClient(APIContextManager): __slots__ = ("_core",) - def __init__(self, session: aiohttp.ClientSession | None = None) -> None: - super().__init__(session) + def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None: + super().__init__(session, user_agent=user_agent) self._core: MegaCore = MegaCore(self._api) if hasattr(sys, "ps1"): setup_logger(logging.DEBUG) diff --git a/src/mega/data_structures.py b/src/mega/data_structures.py index 431ec98..475ed82 100644 --- a/src/mega/data_structures.py +++ b/src/mega/data_structures.py @@ -134,6 +134,9 @@ def dump(self) -> dict[str, Any]: """Get a JSONable dict representation of this object""" return dataclasses.asdict(self) + def __json__(self) -> dict[str, Any]: + return self.dump() + def _shallow_dump(self) -> dict[str, Any]: return {name: getattr(self, name) for name in _fields(type(self))} diff --git a/src/mega/transfer_it.py b/src/mega/transfer_it.py index 5b42c52..a0bb1da 100644 --- a/src/mega/transfer_it.py +++ b/src/mega/transfer_it.py @@ -46,8 +46,8 @@ async def post(self, json: dict[str, Any] | list[dict[str, Any]], params: dict[s class TransferItClient(APIContextManager): - def __init__(self, session: aiohttp.ClientSession | None = None) -> None: # pyright: ignore[reportMissingSuperCall] - self._api = TransferItAPI(session) + def __init__(self, session: aiohttp.ClientSession | None = None, *, user_agent: str | None = None) -> None: # pyright: ignore[reportMissingSuperCall] + self._api = TransferItAPI(session, user_agent=user_agent) @property def progress_bar(self) -> _GeneratorContextManager[None, None, None]: