From d478ed0d4c9f4a01653f1f6de853caa2b213a292 Mon Sep 17 00:00:00 2001 From: disqualifier Date: Thu, 2 Jul 2026 23:30:26 -0400 Subject: [PATCH] fix: warn on dropped OAuth token, skip wasted final backoff, honor external cancel, fix digit-run split (v0.1.8) aiomail-7: log a truncated body when a 200 token response has no access_token instead of silently discarding it. aiomail-8: connect() no longer sleeps after the last failed retry attempt. aiomail-9: _discard_mail now distinguishes its own task.cancel() from an external cancellation via asyncio.shield, so connect() no longer resists cancellation. aiomail-10: the digit-run fallback now matches contiguous \d+ runs instead of mashing digits across punctuation within a whitespace token (e.g. a date). aiomail-11: README's dynamic-matching example uses email.utils.parseaddr since a real From header is not a bare address. Also compresses essay docstrings to one-line-plus-nuance with zero behavior change (re-verified against the same negative controls). Signed-off-by: disqualifier --- README.md | 16 +++++--- pyproject.toml | 2 +- src/aiomail/__init__.py | 2 +- src/aiomail/auth.py | 31 ++++----------- src/aiomail/client.py | 87 +++++++++++++++++------------------------ src/aiomail/extract.py | 42 ++++++++------------ src/aiomail/oauth.py | 26 ++++++------ src/aiomail/retrieve.py | 35 ++++++----------- 8 files changed, 97 insertions(+), 144 deletions(-) diff --git a/README.md b/README.md index aae38b3..c4ea736 100644 --- a/README.md +++ b/README.md @@ -11,22 +11,22 @@ This reads codes from email; it does not generate them (that is `pyotp`'s job). `requirements.txt`: ``` -aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.7 +aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.8 # OAuth token providers (Microsoft / Google) need the extra: -aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.7 +aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.8 ``` Direct: ```bash -pip install "aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.7" -pip install "aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.7" +pip install "aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.8" +pip install "aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.8" ``` Requires `aioimaplib` and `beautifulsoup4` (pulled transitively). The `oauth` extra adds `aiohttp` for the refresh-token providers. -Drop the `@v0.1.7` suffix from the line above to install the latest unpinned. +Drop the `@v0.1.8` suffix from the line above to install the latest unpinned. ## Password auth @@ -61,11 +61,15 @@ Credentials are always supplied by you — nothing is hardcoded. ```python import re +from email.utils import parseaddr await retrieve_otp(client, sender="uber.com") # substring await retrieve_otp(client, sender=re.compile(r"no-?reply@.*\.io")) # regex -await retrieve_otp(client, sender=lambda f: f.endswith("@x.com")) # callable +await retrieve_otp(client, sender=lambda f: parseaddr(f)[1].endswith("@x.com")) # callable ``` +A real `From` header is `Name `, not a bare address — `parseaddr` pulls +the address out before matching (a plain `f.endswith(...)` would never match). + Subject headers are RFC2047-decoded before matching/extraction, so providers that encode non-ASCII subjects (`=?utf-8?B?...?=`) still match on plain text. diff --git a/pyproject.toml b/pyproject.toml index 8e09833..a412d49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "aiomail" -version = "0.1.7" +version = "0.1.8" description = "async IMAP one-time-code retrieval with password/OAuth2 auth and dynamic matching" requires-python = ">=3.10" dependencies = [ diff --git a/src/aiomail/__init__.py b/src/aiomail/__init__.py index e2630fc..1368e76 100644 --- a/src/aiomail/__init__.py +++ b/src/aiomail/__init__.py @@ -31,4 +31,4 @@ __all__ = [ "DEFAULT_FOLDERS", ] -__version__ = "0.1.7" +__version__ = "0.1.8" diff --git a/src/aiomail/auth.py b/src/aiomail/auth.py index ef3312a..dc23857 100644 --- a/src/aiomail/auth.py +++ b/src/aiomail/auth.py @@ -1,9 +1,6 @@ """authentication mechanisms for the IMAP client. -an `Auth` is anything that can authenticate an already-connected aioimaplib -client. two ship: `PasswordAuth` (plain LOGIN) and `OAuth2Auth` (XOAUTH2 using an -access token, either passed directly or pulled from a token provider callable). -credentials are always injected by the caller, never hardcoded here. +`PasswordAuth` (LOGIN) and `OAuth2Auth` (XOAUTH2); credentials always injected. """ import base64 import logging @@ -39,11 +36,7 @@ class PasswordAuth: def _as_str(token) -> str: - """coerce a token to str (a provider may hand back bytes) - - both XOAUTH2 entrypoints downstream need a str (one .encode()s it, the SASL - builder interpolates it), so normalize here rather than crashing on bytes. - """ + """coerce a token to str (a provider may hand back bytes); both XOAUTH2 entrypoints downstream need str""" return token.decode() if isinstance(token, bytes) else token @@ -54,12 +47,7 @@ def _sasl_xoauth2(user: str, token: str) -> str: class OAuth2Auth: - """XOAUTH2 auth using an access token or a token provider - - pass `token` for a static token, or `token_provider` (sync or async callable) - to fetch a fresh one at connect time — the provider path is what makes refresh - flows "easy": hand it a provider from aiomail.oauth and forget about it. - """ + """XOAUTH2 auth: a static `token`, or a `token_provider` (sync/async callable) fetched fresh at connect time""" def __init__( self, @@ -85,18 +73,15 @@ class OAuth2Auth: async def authenticate(self, mail) -> None: token = await self._resolve_token() - # aioimaplib's mail.xoauth2(user, token) builds the SASL string by f-string - # interpolating the token, so token MUST be str — passing bytes interpolates - # the b'...' repr and corrupts the Bearer value. _resolve_token already - # returns str (via _as_str). clients lacking .xoauth2 are driven via the - # SASL callback from _sasl_xoauth2. + # mail.xoauth2(user, token) f-string-interpolates token, so it MUST be str + # — bytes would interpolate the b'...' repr and corrupt the Bearer value. + # _resolve_token already guarantees str via _as_str. xoauth2 = getattr(mail, "xoauth2", None) if xoauth2 is not None: result, data = await xoauth2(self.user, token) elif hasattr(mail, "authenticate"): - # escape hatch for a non-aioimaplib client: the shipped aioimaplib IMAP4 - # always has .xoauth2 and never .authenticate, so this branch never runs - # for it; the SASL-callback signature here is untested against any driver + # escape hatch for a non-aioimaplib client (aioimaplib's IMAP4 always + # has .xoauth2, never .authenticate); untested against any real driver result, data = await mail.authenticate( "XOAUTH2", lambda _: _sasl_xoauth2(self.user, token) ) diff --git a/src/aiomail/client.py b/src/aiomail/client.py index 84ccc2a..03cd206 100644 --- a/src/aiomail/client.py +++ b/src/aiomail/client.py @@ -1,19 +1,9 @@ -"""async IMAP client wrapping aioimaplib. +"""async IMAP client wrapping aioimaplib: connect/retry/reconnect/close plus folders/search/fetch/mark-seen. -a thin, provider-agnostic client: it owns the connection lifecycle (connect with -retries, reconnect-on-stale, close) and exposes the handful of operations the OTP -flow needs (folders, search, fetch, mark-seen). auth is injected, so the same -client serves password and OAuth accounts. - -reconnect-on-stale re-selects whatever folder was selected before the drop, so -search/fetch/store keep working against a reconnected session. sequence-number ids -from before a reconnect are not valid afterward (a fresh SELECT can renumber the -mailbox) — pass `use_uid=True` if ids need to survive a reconnect. - -concurrency: one `IMAPClient` instance is not safe to drive from multiple -concurrent tasks/coroutines without external serialization; connect/reconnect -internally uses a lock to avoid corrupting `_mail`, but overlapping calls to the -same instance are not the intended usage pattern. +auth is injected. reconnect-on-stale re-selects the prior folder, but sequence-number ids from +before a reconnect are not valid after (a fresh SELECT can renumber the mailbox) +— pass `use_uid=True` if ids need to survive a reconnect. one instance is not +safe for concurrent callers beyond the internal connect/reconnect lock. """ import asyncio import email @@ -36,13 +26,9 @@ _LIST_RE = re.compile(rb'^\([^)]*\)\s+(?:"[^"]*"|NIL)\s+(.+)$') def _folder_name(raw: bytes) -> Optional[str]: """extract the folder name from a LIST reply line, delimiter-agnostic - parses the real reply form `(flags) "" ` so any server hierarchy - delimiter works (not just "/"); returns None if the line doesn't match the - canonical shape. aioimaplib appends the tagged completion text (e.g. `b"LIST - completed."`) as the final entry of the same response list that carries the - untagged LIST lines — it never matches LIST syntax, so returning None here - (instead of a last-token rsplit fallback) lets the caller drop it instead of - treating it as a phantom folder name. + returns None on a non-matching line instead of a last-token rsplit fallback, + so the tagged completion line aioimaplib appends to the same response list + (e.g. `b"LIST completed."`) gets dropped instead of read as a phantom folder. """ match = _LIST_RE.match(raw.strip()) if not match: @@ -53,9 +39,8 @@ def _folder_name(raw: bytes) -> Optional[str]: class IMAPClient: """connection-managing IMAP client driven by an injected auth mechanism - note: `use_uid` selects UID vs sequence-number addressing and is independent - of `use_ssl` — the two were conflated in an earlier draft (`use_uid = use_ssl`), - which is a bug; they are unrelated concerns. + `use_uid` (UID vs sequence-number addressing) is independent of `use_ssl` — + an earlier draft conflated them (`use_uid = use_ssl`); unrelated concerns. """ def __init__( @@ -90,9 +75,8 @@ class IMAPClient: async def connect(self) -> bool: """open a connection and authenticate, retrying with linear backoff - serialized by an internal lock: concurrent callers on the same instance queue - up rather than tearing down each other's in-progress handshake. a connection - superseded while a caller waited is logged out, never silently overwritten. + serialized by an internal lock: queued callers never tear down each + other's in-progress handshake, and a superseded connection is logged out. """ async with self._lock: return await self._connect_locked() @@ -117,25 +101,31 @@ class IMAPClient: if self._mail is not None: await self._discard_mail(self._mail) self._mail = None - await asyncio.sleep(2 * (attempt + 1)) + if attempt < self.max_retries - 1: + await asyncio.sleep(2 * (attempt + 1)) return False @staticmethod async def _discard_mail(mail) -> None: """tear down a half-built IMAP4 without leaking its connect task - aioimaplib's IMAP4 schedules `create_connection` as a fire-and-forget task it - never retrieves; on a refused connection that task raises and asyncio logs a - noisy "Task exception was never retrieved" traceback. cancel/await it here (and - retrieve its exception) before discarding, so a failed connect stays quiet. + aioimaplib schedules `create_connection` as a fire-and-forget task; on a + refused connection it raises and asyncio logs a noisy "Task exception was + never retrieved" traceback unless retrieved here first. """ task = getattr(mail, "_client_task", None) if task is not None and not task.done(): task.cancel() if task is not None: + # shield distinguishes our own task.cancel() from an external cancel of + # this coroutine: a bare `await task` swallowed both, resisting + # cancellation. under shield, CancelledError here means external only. try: - await task - except (asyncio.CancelledError, Exception): + await asyncio.shield(task) + except asyncio.CancelledError: + if not task.cancelled(): + raise + except Exception: pass try: await mail.logout() @@ -160,15 +150,10 @@ class IMAPClient: async def ensure_connection(self) -> bool: """return a live, SELECTED-if-applicable connection, reconnecting if the link is stale - a reconnect only re-authenticates (state AUTH); if a folder was selected before - the drop, it is re-selected here so search/fetch/store keep working afterward. - note: sequence-number ids from before the reconnect are NOT valid against the - new session (a fresh SELECT can renumber/re-EXISTS the mailbox) unless - use_uid=True, in which case UIDs remain stable across the reconnect. - - serialized by an internal lock, so concurrent callers on the same instance - never race each other's connect/reconnect; a caller that waits behind another - rechecks liveness first instead of tearing down a connection that just came up. + re-selects the previously-selected folder after a reconnect; sequence-number + ids from before the reconnect are NOT valid against the new session (a fresh + SELECT can renumber the mailbox) unless use_uid=True. serialized by the + internal lock so a queued caller rechecks liveness before reconnecting. """ async with self._lock: if self._mail is not None: @@ -249,8 +234,7 @@ class IMAPClient: try: ids.append(int(token)) except (TypeError, ValueError): - # tolerate a malformed/non-numeric token in the SEARCH response - # instead of crashing the whole search + # tolerate a malformed/non-numeric SEARCH token instead of crashing log.debug("skipping non-numeric search token: %r", token) return sorted(set(ids), reverse=True) @@ -270,14 +254,13 @@ class IMAPClient: if result != "OK" or not data: return None for item in data: - # aioimaplib stores the literal message payload as the only bytearray in - # the response; every other line (including the ` FETCH (...` header) - # is plain bytes. select by structure, not length — a length heuristic - # mismatches the header line for any 2+ digit id or a BODY[]/UID fetch. + # select by structure, not length: aioimaplib stores the literal payload + # as the only bytearray in the response, so a length heuristic would + # mismatch the header line for a 2+ digit id or a BODY[]/UID fetch if isinstance(item, bytearray): return email.message_from_bytes(bytes(item)) - # cross-version fallback: aioimaplib 2.0.x never yields tuples here, but an - # imaplib-style (header, payload) tuple is handled if a future/alt driver does + # cross-version fallback for an imaplib-style (header, payload) tuple; + # aioimaplib 2.0.x never yields one here if isinstance(item, tuple) and len(item) > 1: return email.message_from_bytes(item[1]) return None diff --git a/src/aiomail/extract.py b/src/aiomail/extract.py index 731303d..71c432a 100644 --- a/src/aiomail/extract.py +++ b/src/aiomail/extract.py @@ -1,9 +1,8 @@ -"""code extraction and dynamic matching for email messages. +"""code extraction and dynamic matching for email messages, pure logic with no network IO. -this module is pure logic with no network IO, so it is the easiest part to -unit-test. `extract_code` pulls a one-time code out of a message; `as_predicate` -turns a string / compiled regex / callable into a uniform match function used to -filter senders and subjects. +`extract_code` pulls a one-time code out of a message; `as_predicate` turns a +string / compiled regex / callable into a uniform match function for filtering +senders and subjects. """ import email.message import logging @@ -41,12 +40,10 @@ def _compile(patterns: Sequence[Union[str, Pattern]]) -> list[Pattern]: def decode_header_value(raw: str) -> str: - """decode an RFC2047 encoded-word header (=?charset?B/Q?...?=) to text + """decode an RFC2047 encoded-word header (=?charset?B/Q?...?=) to text, falling back to raw on failure - messages are parsed with email.message_from_bytes (compat32 policy), which - leaves encoded-word headers raw; decode before scanning so the subject-first - pass sees real text instead of base64/quoted-printable. falls back to the - raw value if decoding fails. + messages are parsed with the compat32 policy, which leaves encoded-word + headers raw; decode before scanning so matching sees real text. """ if not raw: return raw @@ -92,10 +89,9 @@ def _scan(text: str, patterns: list[Pattern], lengths: set[int]) -> Optional[str m = pat.search(text) if m: return m.group(1) if m.groups() else m.group(0) - for token in re.split(r"\s+", text): - digits = "".join(c for c in token if c.isdigit()) - if digits and len(digits) in lengths: - return digits + for run in re.findall(r"\d+", text): + if len(run) in lengths: + return run return None @@ -107,10 +103,9 @@ def extract_code( ) -> Optional[str]: """extract a one-time code from a message, subject first then body parts - `patterns` are regexes tried in order; the first capturing group wins, or the - whole match if a pattern has no groups. when no pattern matches a block, any - standalone digit run whose length is in `lengths` is returned. both knobs are - parameters so callers can tune per provider without forking this function. + `patterns` are regexes tried in order (first capturing group wins, else the + whole match); if none hit, a standalone digit run whose length is in + `lengths` is returned. both are parameters so callers tune per provider. """ compiled = _compile(patterns) length_set = set(lengths) @@ -128,19 +123,14 @@ def extract_code( def as_predicate(spec: MatchSpec) -> Callable[[Optional[str]], bool]: - """normalize a match spec into a predicate over an optional string - - accepts None (matches everything), a precompiled regex (search), a callable - (used directly), or a plain string (case-insensitive substring). this is what - gives sender/subject filtering its "string or regex, caller's choice" behavior. - """ + """normalize a match spec (None, regex, callable, or substring) into a predicate over an optional string""" if spec is None: return lambda value: True if isinstance(spec, re.Pattern): return lambda value: bool(spec.search(value or "")) if callable(spec): - # coalesce None like the string/regex branches so the documented Optional[str] - # predicate contract holds even if a caller's callable assumes a real string + # coalesce None like the other branches so the Optional[str] contract holds + # even if the caller's callable assumes a real string return lambda value: bool(spec(value or "")) needle = str(spec).lower() return lambda value: needle in (value or "").lower() diff --git a/src/aiomail/oauth.py b/src/aiomail/oauth.py index ab9c55a..b3c1701 100644 --- a/src/aiomail/oauth.py +++ b/src/aiomail/oauth.py @@ -1,10 +1,7 @@ -"""optional OAuth2 token providers (refresh-token -> access-token). +"""optional OAuth2 token providers (refresh-token -> access-token) for `OAuth2Auth`, credentials always caller-supplied. -these turn a refresh token into a fresh access token for `OAuth2Auth`. they need -aiohttp, kept as an optional extra so the core stays light; importing this module -without aiohttp raises a clear error only when a provider is instantiated. - -credentials (client_id, refresh_token) are always supplied by the caller. +aiohttp is an optional extra so the core stays light; missing it raises a clear +error only when a provider is instantiated, not on import. """ import asyncio import logging @@ -76,16 +73,21 @@ class _RefreshTokenProvider: async with aiohttp.ClientSession(timeout=timeout) as session: async with session.post(endpoint, data=data) as resp: if resp.status == 200: - # content_type=None: some token endpoints return a 200 with - # text/plain or text/javascript; default json() would raise - # ContentTypeError and discard a valid token body - token = (await resp.json(content_type=None)).get("access_token") + # content_type=None: some endpoints return 200 as text/plain + # or text/javascript; default json() would raise ContentTypeError + body_json = await resp.json(content_type=None) + token = body_json.get("access_token") if token: self._failures = 0 return token + # log a truncated body (never whole, may carry sensitive + # material) so a 200-with-no-token isn't a silent drop + log.warning( + "token endpoint %s -> 200 with no access_token: %s", + endpoint, str(body_json)[:200], + ) else: - # log a truncated error body only — a token-endpoint - # response can carry sensitive material; never dump it whole + # truncated only — the body may carry sensitive material body = (await resp.text())[:200] log.warning("token endpoint %s -> %s: %s", endpoint, resp.status, body) except Exception as exc: diff --git a/src/aiomail/retrieve.py b/src/aiomail/retrieve.py index a87aa5a..85d0e6a 100644 --- a/src/aiomail/retrieve.py +++ b/src/aiomail/retrieve.py @@ -1,9 +1,7 @@ -"""orchestration: find the most recent valid OTP across folders. +"""orchestration: `retrieve_otp` ties the client and extractor together to find the most recent valid OTP. -`retrieve_otp` ties the client and extractor together. sender/subject accept the -flexible match specs from `extract`, folders/age/patterns are all parameters with -sane defaults, and provider quirks live in the arguments rather than hardcoded -branches inside the function. +sender/subject accept the flexible match specs from `extract`; provider quirks +(folders, age, patterns) live in the arguments, not hardcoded branches. """ import asyncio import logging @@ -29,15 +27,11 @@ DEFAULT_FOLDERS: Sequence[str] = ("INBOX", "Junk", "Spam", "Archive", "All Mail" def _server_query(sender: MatchSpec, subject: MatchSpec, match_field: str = "from") -> str: - """build a narrowing IMAP query from plain-string specs only + """build a narrowing IMAP query from plain-string specs only, falling back to ALL for regex/callable specs - only plain strings translate to server-side filters; regex and callable specs - fall back to ALL and are filtered client-side, so dynamic matching always works - even when the server cannot express it. `match_field` selects which header the - `sender` spec searches: "from" filters by the sender address (default); "to" - filters by the recipient address (the per-user alias the code was sent to) OR - the From header, matching the client-side check's forwarded-From fallback so - the server query never narrows out a result the client would have accepted. + `match_field="to"` searches TO OR FROM (a forwarded code may keep the + original From), matching the client-side forwarded-From fallback so the + server query never narrows out a result the client would have accepted. """ parts: List[str] = [] if isinstance(sender, str): @@ -58,9 +52,8 @@ def _age_seconds(message) -> Optional[float]: try: dt = parsedate_to_datetime(raw) if dt.tzinfo is None: - # parsedate_to_datetime returns a naive datetime for a "-0000" zone - # (RFC 2822: unknown/unspecified offset); treat it as UTC rather - # than letting .timestamp() interpret the wall time as local + # naive datetime means a "-0000" zone (RFC 2822: unspecified offset); + # force UTC rather than letting .timestamp() read it as local dt = dt.replace(tzinfo=timezone.utc) return time.time() - dt.timestamp() except (TypeError, ValueError) as exc: @@ -86,13 +79,9 @@ async def retrieve_otp( ) -> Optional[str]: """return the newest OTP matching the filters, or None - sender/subject accept a substring, a compiled regex, or a callable. `match_field` - selects which header the `sender` spec is matched against: "from" (default) - matches the sender address; "to" matches the recipient address (the per-user - alias the code was sent to) and additionally accepts a forwarded match on the - From header, so a forwarded code still resolves. folders, patterns, code lengths, - max age and retry behavior are all tunable. set `max_age=None` to disable the - freshness check. + sender/subject accept a substring, a compiled regex, or a callable. + `match_field="to"` matches the recipient address and additionally accepts a + forwarded match on From. set `max_age=None` to disable the freshness check. """ folders = list(folders) if folders is not None else list(DEFAULT_FOLDERS) sender_ok = as_predicate(sender)