# Copyright 2011 Chris Siebenmann # Copyright 2024 Paul Arthur MacIain # GNU General Public License v3.0 (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt) import unittest from miltertest import codec from miltertest import constants from miltertest import ( MilterConnection, MilterError, ) class ConvError(Exception): pass # --- # test infrastructure # These are from the perspective of the socket; it expects you to read # or write. READ, WRITE = object(), object() # A fake socket object that implements .recv() and .sendall(). # It is fed a conversation that it expects (a sequence of read and # write operations), and then verifies that the sequence that happens # is what you told it to expect. # Because this is specific to verifying the milter conversation, it # does not bother having to know exactly what the written messages are; # for our purpose, it is enough to know their type. # (Optionally it can also know the result dictionary and verify it, # because that turned out to be necessary.) class FakeSocket: def __init__(self, conv=None): if conv is None: conv = [] self.conv = conv self.cindex = 0 # verify that a .recv() or .sendall() is proper, ie that it # is the next expected action. def _verify_conv(self, adir): if self.cindex >= len(self.conv): raise ConvError('unexpected action') if adir != self.conv[self.cindex][0]: raise ConvError('sequence mismatch') def _add(self, adir, what): self.conv.append((adir, what)) def addReadMsg(self, cmd, **args): """Add a message to be read; arguments are as per encode_msg.""" self._add(READ, codec.encode_msg(cmd, **args)) def addRead(self, buf): """Add a raw string to be read.""" self._add(READ, buf) def addWrite(self, cmd): """Add an expected write command.""" self._add(WRITE, (cmd,)) def addFullWrite(self, cmd, **args): """Add an expected write command and its full parameters.""" self._add(WRITE, (cmd, args)) def addMTAWrite(self, cmd): """Add an expected write command and then an SMFIR_CONTINUE reply to it.""" self._add(WRITE, (cmd,)) self.addReadMsg(constants.SMFIR_CONTINUE) def isEmpty(self): """Returns whether or not all expected reads and writes have been consumed.""" return self.cindex == len(self.conv) # # The actual socket routines we emulate. def recv(self, nbytes): self._verify_conv(READ) # nbytes should be at least as large as what we are # scheduled to send. _, obj = self.conv[self.cindex] self.cindex += 1 if isinstance(obj, (list, tuple)): obj = codec.encode_msg(obj[0], **obj[1]) if len(obj) > nbytes: raise ConvError('short read') return obj def sendall(self, buf): self._verify_conv(WRITE) # We verify that we got the right sort of stuff r = codec.decode_msg(buf) _, wres = self.conv[self.cindex] self.cindex += 1 otype = wres[0] if r[0] != otype: raise ConvError(f'Received unexpected reply {r[0]} instead of {otype}') if len(wres) > 1 and r[1] != wres[1]: raise ConvError(f'Unexpected reply parameters: {r[1]} instead of {wres[1]}') # ----- # class basicTests(unittest.TestCase): def testShortReads(self): """Test that we correctly read multiple times to reassemble a short message, and that we get the right answer.""" ams = constants.SMFIC_CONNECT adict = {'hostname': 'localhost', 'family': '4', 'port': 1678, 'address': '127.10.10.1'} msg = codec.encode_msg(ams, **adict) msg1, msg2 = msg[:10], msg[10:] s = FakeSocket() s.addRead(msg1) s.addRead(msg2) mbuf = MilterConnection(s) rcmd, rdict = mbuf._recv() self.assertEqual(ams, rcmd) self.assertEqual(adict, rdict) self.assertTrue(s.isEmpty()) def testProgressReads(self): """Test that we correctly read multiple progress messages before getting the real one.""" s = FakeSocket() s.addReadMsg(constants.SMFIR_PROGRESS) s.addReadMsg(constants.SMFIR_PROGRESS) s.addReadMsg(constants.SMFIR_PROGRESS) s.addReadMsg( constants.SMFIR_DELRCPT, rcpt=[ '', ], ) mbuf = MilterConnection(s) rcmd, rdict = mbuf.recv() self.assertEqual(rcmd, constants.SMFIR_DELRCPT) self.assertTrue(s.isEmpty()) class continuedTests(unittest.TestCase): def testHeaders(self): """Test that we handle writing a sequence of headers in the way that we expect.""" s = FakeSocket() hdrs = (('From', 'Chris'), ('To', 'Simon'), ('Subject', 'Yak')) for _ in hdrs: s.addMTAWrite(constants.SMFIC_HEADER) mbuf = MilterConnection(s) rcmd, rdict = mbuf.send_headers(hdrs) self.assertEqual(rcmd, constants.SMFIR_CONTINUE) self.assertTrue(s.isEmpty()) def testShortHeaders(self): """Test that we return early from a series of header writes if SMFIR_CONTINUE is not the code returned.""" s = FakeSocket() hdrs = (('From', 'Chris'), ('To', 'Simon'), ('Subject', 'Yak')) s.addMTAWrite(constants.SMFIC_HEADER) s.addWrite(constants.SMFIC_HEADER) s.addReadMsg(constants.SMFIR_ACCEPT) with self.assertRaises(MilterError): MilterConnection(s).send_headers(hdrs) self.assertTrue(s.isEmpty()) def testBodySequence(self): """Test that we handle writing a large body in the way we expect.""" s = FakeSocket() body = 3 * 65535 * '*' s.addMTAWrite(constants.SMFIC_BODY) s.addMTAWrite(constants.SMFIC_BODY) s.addMTAWrite(constants.SMFIC_BODY) mbuf = MilterConnection(s) rcmd, rdict = mbuf.send_body(body) self.assertEqual(rcmd, constants.SMFIR_CONTINUE) self.assertTrue(s.isEmpty()) def testShortBody(self): """Test that we return early from a series of body writes if SMFIR_CONTINUE is not the code returned.""" s = FakeSocket() body = 3 * 65535 * '*' s.addMTAWrite(constants.SMFIC_BODY) s.addWrite(constants.SMFIC_BODY) s.addReadMsg(constants.SMFIR_ACCEPT) with self.assertRaises(MilterError): MilterConnection(s).send_body(body) self.assertTrue(s.isEmpty()) optneg_mta_pairs = ( ((constants.SMFI_V6_ACTS, constants.SMFI_V6_PROT), (constants.SMFI_V6_ACTS, constants.SMFI_V6_PROT)), ((constants.SMFI_V2_ACTS, constants.SMFI_V2_PROT), (constants.SMFI_V2_ACTS, constants.SMFI_V2_PROT)), ((0x10, 0x10), (0x10, 0x10)), ) def testMTAOptneg(self): """Test that the MTA version of option negotiation returns what we expect it to.""" for a, b in self.optneg_mta_pairs: s = FakeSocket() s.addWrite(constants.SMFIC_OPTNEG) s.addReadMsg(constants.SMFIC_OPTNEG, version=constants.MILTER_VERSION, actions=a[0], protocol=a[1]) # strict=True would blow up on the last test. ract, rprot = MilterConnection(s).optneg_mta(strict=False) self.assertEqual(ract, b[0]) self.assertEqual(rprot, b[1]) optneg_exc_errors = ( (constants.SMFI_V1_ACTS, 0xFFFFFF), (0xFFFFFF, constants.SMFI_V1_PROT), (0xFFFFFF, 0xFFFFFF), ) def testMilterONOutside(self): """Test that the MTA version of option negotiation errors out if there are excess bits in the milter reply.""" for act, prot in self.optneg_exc_errors: s = FakeSocket() s.addWrite(constants.SMFIC_OPTNEG) s.addReadMsg(constants.SMFIC_OPTNEG, version=constants.MILTER_VERSION, actions=act, protocol=prot) bm = MilterConnection(s) self.assertRaises(MilterError, bm.optneg_mta) optneg_milter_pairs = ( # The basic case; MTA says all V6 actions, all V6 protocol # exclusions, we say we'll take all actions and we want the # MTA not to exclude any protocol steps. ((constants.SMFI_V6_ACTS, constants.SMFI_V6_PROT), (constants.SMFI_V6_ACTS, 0x0)), # MTA offers additional protocol exclusions, we tell it not # to do them but to do all V6 protocol actions. ((constants.SMFI_V6_ACTS, 0xFFFFFF), (constants.SMFI_V6_ACTS, 0xE00000)), # MTA offers additional actions, we decline. ((constants.SMFI_V6_ACTS, constants.SMFI_V6_PROT), (constants.SMFI_V2_ACTS, 0x00)), ) def testMilterOptneg(self): """Test the milter version of option negotiation.""" for a, b in self.optneg_milter_pairs: s = FakeSocket() s.addReadMsg(constants.SMFIC_OPTNEG, version=constants.MILTER_VERSION, actions=a[0], protocol=a[1]) s.addFullWrite(constants.SMFIC_OPTNEG, version=constants.MILTER_VERSION, actions=b[0], protocol=b[1]) ract, rprot = MilterConnection(s).optneg_milter(actions=b[0]) self.assertEqual(ract, b[0]) self.assertEqual(rprot, b[1]) if __name__ == '__main__': unittest.main()