67 lines
3.0 KiB
Python
67 lines
3.0 KiB
Python
"""旧连接撤销探针必须区分新登录拒绝与存量连接失效。"""
|
|
|
|
import imaplib
|
|
import poplib
|
|
import smtplib
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
|
|
import runtime_probe
|
|
|
|
|
|
class RevocationProbeTests(unittest.TestCase):
|
|
def run_probe(self, revoked, custom=False):
|
|
smtp, new_smtp = MagicMock(), MagicMock()
|
|
imap, new_imap = MagicMock(), MagicMock()
|
|
pop, new_pop = MagicMock(), MagicMock()
|
|
client = MagicMock()
|
|
client.success.return_value = {'account': runtime_probe.ACCOUNT}
|
|
smtp.auth.return_value = (235, b'OK')
|
|
smtp.mail.side_effect = [(250, b'OK'), (535 if revoked else 250, b'result')]
|
|
smtp.rcpt.return_value = smtp.rset.return_value = (250, b'OK')
|
|
imap.login.return_value = imap.select.return_value = ('OK', [])
|
|
imap.uid.side_effect = [('OK', [(b'1', b'mail')]), ('NO' if revoked else 'OK', [])]
|
|
pop.list.return_value = (b'+OK', [b'1 123'], 0)
|
|
pop.retr.side_effect = [(b'+OK', [], 0), poplib.error_proto(b'-ERR') if revoked else (b'+OK', [], 0)]
|
|
new_smtp.auth.side_effect = smtplib.SMTPAuthenticationError(535, b'denied')
|
|
new_imap.login.side_effect = imaplib.IMAP4.error('AUTHENTICATIONFAILED')
|
|
new_pop.pass_.side_effect = poplib.error_proto(b'-ERR')
|
|
revoke = MagicMock() if custom else None
|
|
with patch.object(runtime_probe, 'smtp_connect', side_effect=[smtp, new_smtp]), \
|
|
patch.object(runtime_probe, 'LocalIMAP', side_effect=[imap, new_imap]), \
|
|
patch.object(runtime_probe, 'LocalPOP', side_effect=[pop, new_pop]):
|
|
result = runtime_probe.verify_protocol_revocation(client, None, dict.fromkeys(('smtp', 'imap', 'pop3'), 'test-code'), revoke)
|
|
if custom:
|
|
revoke.assert_called_once()
|
|
self.assertEqual(client.success.call_count, 1)
|
|
for protocol in ('smtp', 'imap', 'pop3'):
|
|
self.assertTrue(result[protocol + '_new_login_denied'])
|
|
self.assertEqual(result[protocol + '_existing_denied'], revoked)
|
|
smtp.close.assert_called_once()
|
|
imap.shutdown.assert_called_once()
|
|
pop.close.assert_called_once()
|
|
|
|
def test_new_login_rejection_does_not_hide_old_connection_access(self):
|
|
self.run_probe(False)
|
|
|
|
def test_old_and_new_connections_both_rejected(self):
|
|
self.run_probe(True)
|
|
|
|
def test_client_revocation_uses_supplied_action(self):
|
|
self.run_probe(True, custom=True)
|
|
|
|
def test_partial_connect_failure_closes_prior_connections(self):
|
|
smtp = MagicMock()
|
|
with patch.object(runtime_probe, 'smtp_connect', return_value=smtp), \
|
|
patch.object(runtime_probe, 'LocalIMAP', side_effect=OSError('connection failed')):
|
|
with self.assertRaises(OSError):
|
|
runtime_probe.verify_protocol_revocation(MagicMock(), None, {})
|
|
smtp.close.assert_called_once()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|