优化 AWS 邮件验证码过滤与有效字段输出

aws-email-check/tests/test_imap_id.py

@@ -0,0 +1,143 @@
+"""Offline protocol tests for the IMAP ID handshake."""
+
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import unittest
+from unittest.mock import patch
+from zoneinfo import ZoneInfo
+
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
+import check_mail
+
+
+class FakeIMAP:
+    def __init__(self, id_status="OK", id_data=None, select_status="OK", select_data=None,
+                 id_error=None, login_error=None, login_status="OK", login_data=None):
+        self.calls = []
+        self.id_status = id_status
+        self.id_data = id_data if id_data is not None else [b"ID completed"]
+        self.select_status = select_status
+        self.select_data = select_data if select_data is not None else [b"0"]
+        self.id_error = id_error
+        self.login_error = login_error
+        self.login_status = login_status
+        self.login_data = login_data if login_data is not None else [b"logged in"]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_args):
+        pass
+
+    def xatom(self, name, arguments):
+        self.calls.append(("ID", name, arguments))
+        if self.id_error:
+            raise self.id_error
+        return self.id_status, self.id_data
+
+    def login(self, username, authorization_code):
+        self.calls.append(("LOGIN", username, authorization_code))
+        if self.login_error:
+            raise self.login_error
+        return self.login_status, self.login_data
+
+    def select(self, mailbox, readonly=False):
+        self.calls.append(("SELECT", mailbox, readonly))
+        return self.select_status, self.select_data
+
+    def response(self, name):
+        self.calls.append(("RESPONSE", name))
+        return "OK", [b"123"]
+
+    def uid(self, command, *args):
+        self.calls.append(("UID", command, args))
+        return "OK", [b""]
+
+
+class ImapIdTests(unittest.TestCase):
+    def setUp(self):
+        self.settings = {
+            "host": "imap.163.com",
+            "username": "reader@example.com",
+            "authorization_code": "PRIVATE-AUTH-CODE",
+        }
+        self.zone = ZoneInfo("Asia/Shanghai")
+        self.start, self.end = check_mail.day_window(self.zone, "2026-09-28")
+
+    def run_check(self, fake, settings=None):
+        with TemporaryDirectory() as temp, patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=fake):
+            check_mail._check_imap_unlocked(
+                {"imap": settings or self.settings}, self.zone, self.start, self.end,
+                Path(temp) / "processed.json",
+            )
+
+    def test_163_sends_id_before_login_and_select(self):
+        fake = FakeIMAP()
+        self.run_check(fake)
+        self.assertEqual([call[0] for call in fake.calls[:3]], ["ID", "LOGIN", "SELECT"])
+        self.assertEqual(fake.calls[0][1], "ID")
+        self.assertIn('"name" "IMAPClient"', fake.calls[0][2])
+
+    def test_unknown_host_skips_id_by_default(self):
+        fake = FakeIMAP()
+        self.run_check(fake, {**self.settings, "host": "imap.example.com"})
+        self.assertEqual([call[0] for call in fake.calls[:2]], ["LOGIN", "SELECT"])
+
+    def test_explicit_toggle_overrides_host_default(self):
+        fake = FakeIMAP()
+        self.run_check(fake, {**self.settings, "send_id": False})
+        self.assertEqual(fake.calls[0][0], "LOGIN")
+        fake = FakeIMAP()
+        self.run_check(fake, {**self.settings, "host": "imap.example.com", "send_id": True})
+        self.assertEqual(fake.calls[0][0], "ID")
+
+    def test_id_no_aborts_before_login_with_sanitized_reason(self):
+        fake = FakeIMAP(id_status="NO", id_data=[b"Unsafe Login PRIVATE-AUTH-CODE"])
+        with self.assertRaisesRegex(RuntimeError, "Unsafe Login") as error:
+            self.run_check(fake)
+        self.assertNotIn("PRIVATE-AUTH-CODE", str(error.exception))
+        self.assertEqual([call[0] for call in fake.calls], ["ID"])
+
+    def test_id_exception_is_reported_without_credentials(self):
+        fake = FakeIMAP(id_error=OSError("socket failure PRIVATE-AUTH-CODE"))
+        with self.assertRaisesRegex(RuntimeError, "IMAP ID 命令失败") as error:
+            self.run_check(fake)
+        self.assertNotIn("PRIVATE-AUTH-CODE", str(error.exception))
+        self.assertEqual([call[0] for call in fake.calls], ["ID"])
+
+    def test_select_no_preserves_safe_server_reason(self):
+        fake = FakeIMAP(select_status="NO", select_data=[b"SELECT Unsafe Login PRIVATE-AUTH-CODE"])
+        with self.assertRaisesRegex(RuntimeError, "SELECT Unsafe Login") as error:
+            self.run_check(fake)
+        self.assertNotIn("PRIVATE-AUTH-CODE", str(error.exception))
+        self.assertEqual([call[0] for call in fake.calls[:3]], ["ID", "LOGIN", "SELECT"])
+
+    def test_invalid_send_id_mode_rejected(self):
+        fake = FakeIMAP()
+        with self.assertRaisesRegex(ValueError, "send_id"):
+            self.run_check(fake, {**self.settings, "send_id": "always"})
+        self.assertEqual(fake.calls, [])
+
+    def test_login_exception_cannot_echo_credentials(self):
+        fake = FakeIMAP(login_error=OSError(
+            "authentication failed for reader@example.com: PRIVATE-AUTH-CODE"
+        ))
+        with self.assertRaisesRegex(RuntimeError, "IMAP 登录失败") as error:
+            self.run_check(fake)
+        self.assertNotIn(self.settings["username"], str(error.exception))
+        self.assertNotIn(self.settings["authorization_code"], str(error.exception))
+
+    def test_login_no_cannot_echo_credentials(self):
+        fake = FakeIMAP(login_status="NO", login_data=[
+            b"reader@example.com PRIVATE-AUTH-CODE rejected"
+        ])
+        with self.assertRaisesRegex(RuntimeError, "IMAP 登录失败") as error:
+            self.run_check(fake)
+        self.assertNotIn(self.settings["username"], str(error.exception))
+        self.assertNotIn(self.settings["authorization_code"], str(error.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()