242 lines
12 KiB
Python
242 lines
12 KiB
Python
"""Offline tests for explicit recovery of the last generated report."""
|
||
|
||
from contextlib import redirect_stdout
|
||
from email.message import EmailMessage
|
||
from io import StringIO
|
||
import json
|
||
from pathlib import Path
|
||
import sys
|
||
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, message, uid=b"42"):
|
||
self.message = message
|
||
self.uid_value = uid
|
||
self.fetches = 0
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *_args):
|
||
pass
|
||
|
||
def login(self, *_args):
|
||
return "OK", [b"logged in"]
|
||
|
||
def select(self, *_args, **_kwargs):
|
||
return "OK", [b"1"]
|
||
|
||
def response(self, name):
|
||
assert name == "UIDVALIDITY"
|
||
return "OK", [b"123"]
|
||
|
||
def uid(self, command, *_args):
|
||
if command == "SEARCH":
|
||
return "OK", [self.uid_value]
|
||
if command == "FETCH":
|
||
self.fetches += 1
|
||
return "OK", [
|
||
(b'1 (UID 42 INTERNALDATE "28-Sep-2026 10:00:00 +0800" BODY[] {1}', self.message),
|
||
b")",
|
||
]
|
||
raise AssertionError(command)
|
||
|
||
|
||
class ReplayTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.temp = TemporaryDirectory()
|
||
self.addCleanup(self.temp.cleanup)
|
||
self.root = Path(self.temp.name)
|
||
self.config = self.root / "config.json"
|
||
self.state = self.root / "state" / "processed.json"
|
||
self.config.write_text(json.dumps({
|
||
"imap": {
|
||
"host": "imap.example.com",
|
||
"username": "reader@example.com",
|
||
"authorization_code": "test-only",
|
||
},
|
||
"timezone": "Asia/Shanghai",
|
||
"state_file": "state/processed.json",
|
||
}), encoding="utf-8")
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "AWS Account on Hold: Response Required"
|
||
mail["Message-ID"] = "<report-42@example.com>"
|
||
mail.set_content("AWS Account ID: 123456789012\nYour account has been placed on hold.")
|
||
self.fake = FakeIMAP(mail.as_bytes())
|
||
# Keep the fake INTERNALDATE within the test's day without changing production filtering.
|
||
window = check_mail.day_window(ZoneInfo("Asia/Shanghai"), "2026-09-28")
|
||
fixed_day = patch.object(check_mail, "day_window", return_value=window)
|
||
fixed_day.start()
|
||
self.addCleanup(fixed_day.stop)
|
||
|
||
def run_main(self, *args):
|
||
output = StringIO()
|
||
with redirect_stdout(output):
|
||
check_mail.main(["--config", str(self.config), *args])
|
||
return output.getvalue()
|
||
|
||
def test_dedup_and_explicit_replay_without_connection(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=self.fake):
|
||
first = self.run_main()
|
||
second = self.run_main()
|
||
self.assertIn("【AWS账号风控邮件】", first)
|
||
self.assertIn("AWS账号ID：123456789012", first)
|
||
self.assertEqual(second, check_mail.NO_NEW + "\n")
|
||
self.assertEqual(self.fake.fetches, 1)
|
||
before = self.state.read_bytes()
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", side_effect=AssertionError("unexpected connection")):
|
||
replay = self.run_main("--replay-last")
|
||
self.assertEqual(replay, first)
|
||
self.assertEqual(self.state.read_bytes(), before)
|
||
self.assertEqual(json.loads(before)["last_report"], first.removesuffix("\n"))
|
||
|
||
def test_empty_scan_does_not_replace_previous_report(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=self.fake):
|
||
first = self.run_main()
|
||
empty = FakeIMAP(self.fake.message, uid=b"")
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=empty):
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
self.assertEqual(self.run_main("--replay-last"), first)
|
||
|
||
def test_older_batch_survives_next_report_and_empty_scan(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=self.fake):
|
||
first = self.run_main()
|
||
newer_mail = EmailMessage()
|
||
newer_mail["Subject"] = "AWS Account on Hold: Second Notice"
|
||
newer_mail["Message-ID"] = "<report-43@example.com>"
|
||
newer_mail.set_content("AWS Account ID: 999999999999\nYour account has been placed on hold.")
|
||
newer = FakeIMAP(newer_mail.as_bytes(), uid=b"43")
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=newer):
|
||
second = self.run_main()
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
state_before = self.state.read_bytes()
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", side_effect=AssertionError("unexpected connection")):
|
||
ids = self.run_main("--list-report-batches").splitlines()
|
||
self.assertEqual(len(ids), 2)
|
||
self.assertEqual(self.run_main("--replay-id", ids[1]), first)
|
||
self.assertEqual(self.run_main("--replay-id", ids[0]), second)
|
||
self.assertEqual(self.run_main("--replay-last"), second)
|
||
self.assertEqual(self.state.read_bytes(), state_before)
|
||
|
||
def test_history_is_bounded(self):
|
||
state = {"mailboxes": {}}
|
||
for number in range(check_mail.REPORT_HISTORY_LIMIT + 4):
|
||
check_mail.save_report_batch(state, f"report {number}")
|
||
history = check_mail.report_history(state)
|
||
self.assertEqual(len(history), check_mail.REPORT_HISTORY_LIMIT)
|
||
self.assertEqual(history[0]["report"], "report 4")
|
||
self.assertEqual(history[-1]["report"], f"report {check_mail.REPORT_HISTORY_LIMIT + 3}")
|
||
self.assertEqual(len({item["id"] for item in history}), len(history))
|
||
|
||
def test_legacy_last_report_remains_replayable(self):
|
||
self.state.parent.mkdir(parents=True, exist_ok=True)
|
||
self.state.write_text(json.dumps({"mailboxes": {}, "last_report": "旧报告"}), encoding="utf-8")
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", side_effect=AssertionError("unexpected connection")):
|
||
self.assertEqual(self.run_main("--list-report-batches"), "legacy\n")
|
||
self.assertEqual(self.run_main("--replay-id", "legacy"), "旧报告\n")
|
||
self.assertEqual(self.run_main("--replay-last"), "旧报告\n")
|
||
|
||
def test_missing_batch_and_invalid_history_are_clear_errors(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", side_effect=AssertionError("unexpected connection")):
|
||
with self.assertRaisesRegex(RuntimeError, "报告批次 ID 不存在"):
|
||
self.run_main("--replay-id", "missing")
|
||
self.state.parent.mkdir(parents=True, exist_ok=True)
|
||
self.state.write_text('{"mailboxes": {}, "report_history": [42]}', encoding="utf-8")
|
||
with self.assertRaisesRegex(ValueError, "报告历史格式无效"):
|
||
self.run_main("--list-report-batches")
|
||
|
||
def test_no_saved_report_has_clear_error(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", side_effect=AssertionError("unexpected connection")):
|
||
with self.assertRaisesRegex(RuntimeError, "没有可重放的上次报告"):
|
||
self.run_main("--replay-last")
|
||
self.state.parent.mkdir(parents=True, exist_ok=True)
|
||
self.state.write_text('{"mailboxes": {}}', encoding="utf-8")
|
||
with self.assertRaisesRegex(RuntimeError, "没有可重放的上次报告"):
|
||
self.run_main("--replay-last")
|
||
|
||
def test_local_eml_mode_does_not_write_state(self):
|
||
folder = self.root / "samples"
|
||
folder.mkdir()
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "AWS Account on Hold: Response Required"
|
||
mail["Date"] = "Mon, 28 Sep 2026 10:00:00 +0800"
|
||
mail.set_content("AWS Account ID: 123456789012\nYour account has been placed on hold.")
|
||
(folder / "sample.eml").write_bytes(mail.as_bytes())
|
||
self.assertIn("【AWS账号风控邮件】", self.run_main("--eml-dir", str(folder), "--date", "2026-09-28"))
|
||
self.assertFalse(self.state.exists())
|
||
|
||
def test_filtered_code_is_deduplicated_without_report_history(self):
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "AWS 验证您的身份"
|
||
mail["Message-ID"] = "<code-42@example.com>"
|
||
mail.set_content("AWS 验证码 123456，用于登录。如果发现异常活动，请启用 MFA。")
|
||
fake = FakeIMAP(mail.as_bytes())
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=fake):
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
self.assertEqual(fake.fetches, 1)
|
||
state = json.loads(self.state.read_text(encoding="utf-8"))
|
||
entry = next(iter(state["mailboxes"].values()))
|
||
self.assertEqual(entry["uids"], ["42"])
|
||
self.assertEqual(entry["message_ids"], ["<code-42@example.com>"])
|
||
self.assertNotIn("report_history", state)
|
||
self.assertNotIn("last_report", state)
|
||
|
||
def test_filtered_code_preserves_previous_report_batch(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=self.fake):
|
||
first = self.run_main()
|
||
before = json.loads(self.state.read_text(encoding="utf-8"))
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "AWS verification code"
|
||
mail["Message-ID"] = "<code-43@example.com>"
|
||
mail.set_content("Verification code: 123456. Use this code to sign in.")
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=FakeIMAP(mail.as_bytes(), uid=b"43")):
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
after = json.loads(self.state.read_text(encoding="utf-8"))
|
||
self.assertEqual(after["report_history"], before["report_history"])
|
||
self.assertEqual(self.run_main("--replay-last"), first)
|
||
|
||
def test_filtered_access_validated_notice_is_deduplicated_without_history(self):
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "Fw: Your Request For Accessing AWS Resources Has Been Validated"
|
||
mail["Message-ID"] = "<validated-42@example.com>"
|
||
mail.set_content("Your AWS resource-access request has completed verification.")
|
||
fake = FakeIMAP(mail.as_bytes())
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=fake):
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
self.assertEqual(fake.fetches, 1)
|
||
state = json.loads(self.state.read_text(encoding="utf-8"))
|
||
entry = next(iter(state["mailboxes"].values()))
|
||
self.assertEqual(entry["uids"], ["42"])
|
||
self.assertEqual(entry["message_ids"], ["<validated-42@example.com>"])
|
||
self.assertNotIn("report_history", state)
|
||
self.assertNotIn("last_report", state)
|
||
|
||
def test_filtered_access_validated_notice_preserves_previous_report(self):
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=self.fake):
|
||
first = self.run_main()
|
||
before = json.loads(self.state.read_text(encoding="utf-8"))
|
||
mail = EmailMessage()
|
||
mail["Subject"] = "Fw: 您访问 AWS 资源的请求已得到验证"
|
||
mail["Message-ID"] = "<validated-43@example.com>"
|
||
mail.set_content("<p>现在可以访问 AWS 资源。</p>", subtype="html")
|
||
with patch.object(check_mail.imaplib, "IMAP4_SSL", return_value=FakeIMAP(mail.as_bytes(), uid=b"43")):
|
||
self.assertEqual(self.run_main(), check_mail.NO_NEW + "\n")
|
||
after = json.loads(self.state.read_text(encoding="utf-8"))
|
||
self.assertEqual(after["report_history"], before["report_history"])
|
||
self.assertEqual(after["last_report"], before["last_report"])
|
||
self.assertEqual(self.run_main("--replay-last"), first)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|