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 <dev@disqualifier.me>
This commit is contained in:
2026-07-02 23:30:26 -04:00
parent fd789c2ba2
commit e35918ef2d
8 changed files with 97 additions and 144 deletions
+10 -6
View File
@@ -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 <addr@x.com>`, 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.
+1 -1
View File
@@ -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 = [
+1 -1
View File
@@ -31,4 +31,4 @@ __all__ = [
"DEFAULT_FOLDERS",
]
__version__ = "0.1.7"
__version__ = "0.1.8"
+8 -23
View File
@@ -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)
)
+35 -52
View File
@@ -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) "<delim>" <name>` 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 `<id> 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
+16 -26
View File
@@ -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()
+14 -12
View File
@@ -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:
+12 -23
View File
@@ -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)