diff --git a/src/aiomail/extract.py b/src/aiomail/extract.py index 29f7352..730a20b 100644 --- a/src/aiomail/extract.py +++ b/src/aiomail/extract.py @@ -15,12 +15,13 @@ log = logging.getLogger(__name__) MatchSpec = Union[str, Pattern, Callable[[Optional[str]], bool], None] # patterns tried (in order) against the subject then each body part; the first -# capturing group is the code, or the whole match if there are no groups +# capturing group is the code, or the whole match if there are no groups. no bare +# digit catch-all here - that is the `lengths` fallback, which runs only after every +# text missed every pattern DEFAULT_PATTERNS: Sequence[Union[str, Pattern]] = ( r"(\d{4,8})\s+is\s+your", r"your.{0,20}?code(?:\s+is)?[:\s]+(\d{4,8})", r"code(?:\s+is)?[:\s]+(\d{4,8})", - r"\b(\d{4,8})\b", ) # default digit-run lengths to treat as a standalone code when no pattern hits @@ -79,15 +80,12 @@ def _iter_text(message: email.message.Message) -> Iterator[str]: yield text -def _scan(text: str, patterns: list[Pattern], lengths: set[int]) -> Optional[str]: - """run patterns then a standalone digit-run fallback over one block of text""" +def _match(text: str, patterns: list[Pattern]) -> Optional[str]: + """first pattern hit in one block of text""" for pat in patterns: m = pat.search(text) if m: return m.group(1) if m.groups() else m.group(0) - for run in re.findall(r"\d+", text): - if len(run) in lengths: - return run return None @@ -100,20 +98,21 @@ def extract_code( """extract a one-time code from a message, subject first then body parts `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. + match) over the subject and every body part; only if none hit anywhere is a + standalone digit run whose length is in `lengths` returned, from the same texts. """ compiled = _compile(patterns) length_set = set(lengths) + texts = [decode_header_value(message.get("Subject", "") or ""), *_iter_text(message)] - subject = decode_header_value(message.get("Subject", "") or "") - hit = _scan(subject, compiled, length_set) - if hit: - return hit - - for text in _iter_text(message): - hit = _scan(text, compiled, length_set) + for text in texts: + hit = _match(text, compiled) if hit: return hit + for text in texts: + for run in re.findall(r"\d+", text): + if len(run) in length_set: + return run return None diff --git a/src/aiomail/retrieve.py b/src/aiomail/retrieve.py index ce73eca..5e83914 100644 --- a/src/aiomail/retrieve.py +++ b/src/aiomail/retrieve.py @@ -22,9 +22,12 @@ log = logging.getLogger(__name__) 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, falling back to ALL for regex/callable specs""" - parts: List[str] = [] +def _server_query(sender: MatchSpec, subject: MatchSpec, match_field: str = "from", unseen: bool = False) -> str: + """build a narrowing IMAP query from plain-string specs (ALL when nothing narrows) + + `unseen` adds UNSEEN so a code already consumed (marked seen) is never served again + """ + parts: List[str] = ["UNSEEN"] if unseen else [] if isinstance(sender, str): if match_field == "to": parts.append(f'OR TO "{sender}" FROM "{sender}"') @@ -73,20 +76,20 @@ async def retrieve_otp( 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. + with `mark_seen=True` (default) only UNSEEN messages are searched, so a code is + returned once; `mark_seen=False` searches seen messages too. """ folders = list(folders) if folders is not None else list(DEFAULT_FOLDERS) lengths = list(lengths) sender_ok = as_predicate(sender) subject_ok = as_predicate(subject) - query = _server_query(sender, subject, match_field) + query = _server_query(sender, subject, match_field, unseen=mark_seen) for attempt in range(retries + 1): for folder in folders: if not await client.select(folder): continue ids = await client.search(query) - if not ids and query != "ALL": - ids = await client.search("ALL") if not ids: continue