mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-01 08:42:20 +03:00
Code Issues Packages Projects Releases Wiki Activity

perf(email): fetch headers before body, use UID SEARCH to skip re-fetch

Browse Source 
The IMAP poll loop previously downloaded the entire message body for
every UNSEEN message before running any filter (self-sent, SPF/DKIM,
allow-list), and only learned the UID by parsing it back out of that
fetch response. Rejected messages stay unseen and get re-fetched in
full on every subsequent poll.

Switch to UID SEARCH (UIDs come back directly, no per-message fetch
needed to learn them) so already-processed UIDs are skipped before any
network fetch, then fetch headers only to evaluate every filter — the
full body is downloaded only for messages that pass every check and
are actually delivered. No behavior change: filter outcomes and
\Seen semantics are unchanged.
This commit is contained in:
Till Adam
2026-08-30 17:38:44 +08:00
committed by Xubin Ren
parent 1ac1b35c84
commit f573ecfe56
2 changed files with 113 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/email/runtime.py
+58 -43
View File
@@ -430,7 +430,13 @@ class EmailChannel(BaseChannel):
skipped_uids: set[str], skipped_uids: set[str],
cycle_uids: set[str], cycle_uids: set[str],
) -> list[dict[str, Any]] | None: ) -> list[dict[str, Any]] | None:
"""Fetch messages by arbitrary IMAP search criteria.""" """Fetch messages by arbitrary IMAP search criteria.
Uses UID SEARCH so already-processed UIDs are recognized before any
FETCH at all, then fetches headers only to evaluate every filter — the
full body (and any attachments) is downloaded only for messages that
pass every check and are actually going to be delivered.
"""
mailbox = self.config.imap_mailbox or "INBOX" mailbox = self.config.imap_mailbox or "INBOX"
client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True) client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True)
@@ -438,29 +444,30 @@ class EmailChannel(BaseChannel):
return messages return messages
try: try:
status, data = client.search(None, *search_criteria) status, data = client.uid("SEARCH", None, *search_criteria)
if status != "OK" or not data: if status != "OK" or not data or not data[0]:
return messages return messages
ids = data[0].split() uids = [raw.decode("ascii", errors="ignore") for raw in data[0].split()]
if limit > 0 and len(ids) > limit: if limit > 0 and len(uids) > limit:
ids = ids[-limit:] uids = uids[-limit:]
for imap_id in ids:
status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)") features: _ServerFeatures | None = None
for uid in uids:
if not uid or uid in cycle_uids:
continue
if dedupe and uid in self._processed_uids:
continue
status, fetched = client.uid("FETCH", uid, "(BODY.PEEK[HEADER])")
if status != "OK" or not fetched: if status != "OK" or not fetched:
continue continue
header_bytes = self._extract_message_bytes(fetched)
raw_bytes = self._extract_message_bytes(fetched) if header_bytes is None:
if raw_bytes is None:
continue continue
uid = self._extract_uid(fetched) parsed = BytesParser(policy=policy.default).parsebytes(header_bytes)
if uid and uid in cycle_uids:
continue
if dedupe and uid and uid in self._processed_uids:
continue
parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
sender = parseaddr(parsed.get("From", ""))[1].strip().lower() sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
if not sender: if not sender:
continue continue
@@ -468,9 +475,8 @@ class EmailChannel(BaseChannel):
self.logger.info("From {} ignored: matches bot-owned address", sender) self.logger.info("From {} ignored: matches bot-owned address", sender)
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen") features = self._mark_seen_uid(client, uid, features)
if uid: skipped_uids.add(uid)
skipped_uids.add(uid)
continue continue
# --- Anti-spoofing: verify Authentication-Results --- # --- Anti-spoofing: verify Authentication-Results ---
@@ -482,8 +488,7 @@ class EmailChannel(BaseChannel):
sender, sender,
) )
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if uid: skipped_uids.add(uid)
skipped_uids.add(uid)
continue continue
if self.config.verify_dkim and not dkim_pass: if self.config.verify_dkim and not dkim_pass:
self.logger.warning( self.logger.warning(
@@ -492,18 +497,26 @@ class EmailChannel(BaseChannel):
sender, sender,
) )
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if uid: skipped_uids.add(uid)
skipped_uids.add(uid)
continue continue
if not self.is_allowed(sender): if not self.is_allowed(sender):
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen") features = self._mark_seen_uid(client, uid, features)
if uid: skipped_uids.add(uid)
skipped_uids.add(uid)
continue continue
# Passed every filter — only now fetch the full message body
# (and any attachments) for the message we're actually delivering.
status, full_fetched = client.uid("FETCH", uid, "(BODY.PEEK[])")
if status != "OK" or not full_fetched:
continue
raw_bytes = self._extract_message_bytes(full_fetched)
if raw_bytes is None:
continue
parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
subject = self._decode_header_value(parsed.get("Subject", "")) subject = self._decode_header_value(parsed.get("Subject", ""))
date_value = parsed.get("Date", "") date_value = parsed.get("Date", "")
message_id = parsed.get("Message-ID", "").strip() message_id = parsed.get("Message-ID", "").strip()
@@ -556,10 +569,19 @@ class EmailChannel(BaseChannel):
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen") features = self._mark_seen_uid(client, uid, features)
finally: finally:
self._close_imap_client(client) self._close_imap_client(client)
def _mark_seen_uid(
self, client: Any, uid: str, features: _ServerFeatures | None
) -> _ServerFeatures:
"""Mark a single UID \\Seen, reusing session-learned STORE support."""
if features is None:
features = self._server_features(client)
self._uid_store_flag(client, uid, "\\Seen", features)
return features
def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None: def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None:
if self.config.imap_use_ssl: if self.config.imap_use_ssl:
client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port) client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
@@ -714,11 +736,14 @@ class EmailChannel(BaseChannel):
return data[0].split()[0] return data[0].split()[0]
def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool: def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool:
return self._uid_store_flag(client, uid, "\\Deleted", features)
def _uid_store_flag(self, client: Any, uid: str, flag: str, features: _ServerFeatures) -> bool:
# Optimistic path: try UID STORE first because UID is stable and avoids # Optimistic path: try UID STORE first because UID is stable and avoids
# sequence-number lookup. If this fails once for the session, remember it # sequence-number lookup. If this fails once for the session, remember it
# and use the sequence STORE fallback directly for remaining UIDs. # and use the sequence STORE fallback directly for remaining UIDs.
if features.uid_store is not False: if features.uid_store is not False:
status, _ = client.uid("STORE", uid, "+FLAGS", "(\\Deleted)") status, _ = client.uid("STORE", uid, "+FLAGS", f"({flag})")
if status == "OK": if status == "OK":
features.uid_store = True features.uid_store = True
return True return True
@@ -728,12 +753,12 @@ class EmailChannel(BaseChannel):
# unreliable: resolve the current sequence number from UID and use STORE. # unreliable: resolve the current sequence number from UID and use STORE.
imap_id = self._lookup_imap_id_by_uid(client, uid) imap_id = self._lookup_imap_id_by_uid(client, uid)
if not imap_id: if not imap_id:
self.logger.warning("Post-action skipped: UID {} not found", uid) self.logger.warning("Could not locate UID {} to set flag {}", uid, flag)
return False return False
status, _ = client.store(imap_id, "+FLAGS", "\\Deleted") status, _ = client.store(imap_id, "+FLAGS", flag)
if status != "OK": if status != "OK":
self.logger.warning("Post-action failed: could not mark UID {} as deleted", uid) self.logger.warning("Failed to set flag {} on UID {}", flag, uid)
return False return False
return True return True
@@ -773,16 +798,6 @@ class EmailChannel(BaseChannel):
return bytes(fetched_item[1]) return bytes(fetched_item[1])
return None return None
@staticmethod
def _extract_uid(fetched: list[Any]) -> str:
for item in fetched:
if isinstance(item, tuple) and item and isinstance(item[0], (bytes, bytearray)):
head = bytes(item[0]).decode("utf-8", errors="ignore")
m = re.search(r"UID\s+(\d+)", head)
if m:
return m.group(1)
return ""
@staticmethod @staticmethod
def _decode_header_value(value: str) -> str: def _decode_header_value(value: str) -> str:
if not value: if not value:
nanobot/channels/email/tests/test_email_channel.py
+55 -155
View File
@@ -53,30 +53,7 @@ def _make_raw_email(
def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None: def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
raw = _make_raw_email(subject="Invoice", body="Please pay") raw = _make_raw_email(subject="Invoice", body="Please pay")
class FakeIMAP: fake = _make_fake_imap(raw, uid=b"123")
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(), MessageBus()) channel = EmailChannel(_make_config(), MessageBus())
@@ -86,7 +63,7 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
assert items[0]["sender"] == "alice@example.com" assert items[0]["sender"] == "alice@example.com"
assert items[0]["subject"] == "Invoice" assert items[0]["subject"] == "Invoice"
assert "Please pay" in items[0]["content"] assert "Please pay" in items[0]["content"]
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")] assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
assert skipped_uids == set() assert skipped_uids == set()
# Same UID should be deduped in-process. # Same UID should be deduped in-process.
@@ -98,26 +75,8 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> None: def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> None:
raw = _make_raw_email(subject="Invoice", body="Please pay") raw = _make_raw_email(subject="Invoice", body="Please pay")
class FakeIMAP: fake = _make_fake_imap(raw, uid=b"123")
def login(self, _user: str, _pw: str): monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, _imap_id: bytes, _op: str, _flags: str):
return "OK", [b""]
def logout(self):
return "BYE", [b""]
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
channel = EmailChannel(_make_config(post_action="delete"), MessageBus()) channel = EmailChannel(_make_config(post_action="delete"), MessageBus())
items, skipped_uids = channel._fetch_new_messages() items, skipped_uids = channel._fetch_new_messages()
@@ -130,26 +89,10 @@ def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> No
def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None: def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test") raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
class FakeIMAP: monkeypatch.setattr(
def login(self, _user: str, _pw: str): "nanobot.channels.email.runtime.imaplib.IMAP4_SSL",
return "OK", [b"logged in"] lambda _h, _p: _make_fake_imap(raw, uid=b"123"),
)
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, _imap_id: bytes, _op: str, _flags: str):
return "OK", [b""]
def logout(self):
return "BYE", [b""]
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
channel_skip = EmailChannel( channel_skip = EmailChannel(
_make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True), _make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True),
@@ -545,30 +488,7 @@ async def test_start_keeps_post_actions_for_successful_emails_when_later_deliver
def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) -> None: def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) -> None:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test") raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
class FakeIMAP: fake = _make_fake_imap(raw, uid=b"123")
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus()) channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus())
@@ -576,7 +496,7 @@ def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) ->
assert items == [] assert items == []
assert skipped_uids == {"123"} assert skipped_uids == {"123"}
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")] assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
# Same UID should still be deduped after being ignored. # Same UID should still be deduped after being ignored.
items_again, skipped_again = channel._fetch_new_messages() items_again, skipped_again = channel._fetch_new_messages()
@@ -614,37 +534,14 @@ def test_fetch_new_messages_skips_self_sent_across_identity_sources(
imap_username matches, and must be case-insensitive.""" imap_username matches, and must be case-insensitive."""
raw = _make_raw_email(from_addr=from_header, subject="Loop test") raw = _make_raw_email(from_addr=from_header, subject="Loop test")
class FakeIMAP: fake = _make_fake_imap(raw, uid=b"123")
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(**config_override), MessageBus()) channel = EmailChannel(_make_config(**config_override), MessageBus())
items, _ = channel._fetch_new_messages() items, _ = channel._fetch_new_messages()
assert items == [] assert items == []
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")] assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None: def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None:
@@ -662,15 +559,16 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def search(self, *_args): def uid(self, command: str, *args):
self.search_calls += 1 if command == "SEARCH":
if fail_once["pending"]: self.search_calls += 1
fail_once["pending"] = False if fail_once["pending"]:
raise imaplib.IMAP4.abort("socket error") fail_once["pending"] = False
return "OK", [b"1"] raise imaplib.IMAP4.abort("socket error")
return "OK", [b"123"]
def fetch(self, _imap_id: bytes, _parts: str): if command == "FETCH":
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"] return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -700,10 +598,7 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypatch) -> None: def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypatch) -> None:
raw_first = _make_raw_email(subject="First", body="First body") raw_first = _make_raw_email(subject="First", body="First body")
raw_second = _make_raw_email(subject="Second", body="Second body") raw_second = _make_raw_email(subject="Second", body="Second body")
mailbox_state = { mailbox_state = {"123": raw_first, "124": raw_second}
b"1": {"uid": b"123", "raw": raw_first, "seen": False},
b"2": {"uid": b"124", "raw": raw_second, "seen": False},
}
fail_once = {"pending": True} fail_once = {"pending": True}
class FlakyIMAP: class FlakyIMAP:
@@ -713,20 +608,18 @@ def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypa
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"2"] return "OK", [b"2"]
def search(self, *_args): def uid(self, command: str, *args):
unseen_ids = [imap_id for imap_id, item in mailbox_state.items() if not item["seen"]] if command == "SEARCH":
return "OK", [b" ".join(unseen_ids)] keys = " ".join(sorted(mailbox_state.keys(), key=int))
return "OK", [keys.encode()]
def fetch(self, imap_id: bytes, _parts: str): if command == "FETCH":
if imap_id == b"2" and fail_once["pending"]: uid = args[0]
fail_once["pending"] = False if uid == "124" and fail_once["pending"]:
raise imaplib.IMAP4.abort("socket error") fail_once["pending"] = False
item = mailbox_state[imap_id] raise imaplib.IMAP4.abort("socket error")
header = b"%s (UID %s BODY[] {200})" % (imap_id, item["uid"]) raw = mailbox_state[uid]
return "OK", [(header, item["raw"]), b")"] header = f"{uid} (UID {uid} BODY[] {{200}})".encode()
return "OK", [(header, raw), b")"]
def store(self, imap_id: bytes, _op: str, _flags: str):
mailbox_state[imap_id]["seen"] = True
return "OK", [b""] return "OK", [b""]
def logout(self): def logout(self):
@@ -1044,12 +937,13 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def search(self, *_args): def uid(self, command: str, *args):
self.search_args = _args if command == "SEARCH":
return "OK", [b"5"] self.search_args = args
return "OK", [b"999"]
def fetch(self, _imap_id: bytes, _parts: str): if command == "FETCH":
return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"] return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -1070,7 +964,7 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
assert len(items) == 1 assert len(items) == 1
assert items[0]["subject"] == "Status" assert items[0]["subject"] == "Status"
# search(None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026") # uid("SEARCH", None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.search_args is not None assert fake.search_args is not None
assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026") assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.store_calls == [] assert fake.store_calls == []
@@ -1080,11 +974,12 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
# Security: Anti-spoofing tests for Authentication-Results verification # Security: Anti-spoofing tests for Authentication-Results verification
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _make_fake_imap(raw: bytes): def _make_fake_imap(raw: bytes, uid: bytes = b"500"):
"""Return a FakeIMAP class pre-loaded with the given raw email.""" """Return a FakeIMAP class pre-loaded with the given raw email."""
class FakeIMAP: class FakeIMAP:
def __init__(self) -> None: def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = [] self.store_calls: list[tuple[bytes, str, str]] = []
self.uid_calls: list[tuple] = []
def login(self, _user: str, _pw: str): def login(self, _user: str, _pw: str):
return "OK", [b"logged in"] return "OK", [b"logged in"]
@@ -1092,11 +987,16 @@ def _make_fake_imap(raw: bytes):
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def search(self, *_args): def capability(self):
return "OK", [b"1"] return "OK", [b"IMAP4rev1"]
def fetch(self, _imap_id: bytes, _parts: str): def uid(self, command: str, *args):
return "OK", [(b"1 (UID 500 BODY[] {200})", raw), b")"] self.uid_calls.append((command, *args))
if command == "SEARCH":
return "OK", [uid]
if command == "FETCH":
return "OK", [(b"1 (UID " + uid + b" BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -1292,7 +1192,7 @@ def test_fetch_new_messages_ignores_unauthorized_sender_before_attachments(monke
assert channel._fetch_new_messages() == ([], {"500"}) assert channel._fetch_new_messages() == ([], {"500"})
assert called["attachments"] is False assert called["attachments"] is False
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")] assert ("STORE", "500", "+FLAGS", "(\\Seen)") in fake.uid_calls
def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None: def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None: