diff --git a/Milter/__init__.py b/Milter/__init__.py index 6e9da640717490b3f5eb344e66fa5196c2db7930..85c5ab0e70d9c9706be0866c5fb0973b55debf74 100755 --- a/Milter/__init__.py +++ b/Milter/__init__.py @@ -48,6 +48,12 @@ OPTIONAL_CALLBACKS = { 'header':(P_NR_HDR,P_NOHDRS) } +MACRO_CALLBACKS = { + 'connect': M_CONNECT, + 'hello': M_HELO, 'envfrom': M_ENVFROM, 'envrcpt': M_ENVRCPT, + 'data': M_DATA, 'eom': M_EOM, 'eoh': M_EOH +} + ## @private R = re.compile(r'%+') @@ -141,6 +147,7 @@ def nocallback(func): except KeyError: raise ValueError( '@nocallback applied to non-optional method: '+func.__name__) + @wraps(func) def wrapper(self,*args): if func(self,*args) != CONTINUE: raise RuntimeError('%s return code must be CONTINUE with @nocallback' @@ -173,6 +180,19 @@ def noreply(func): wrapper.milter_protocol = nr_mask return wrapper +## Function decorator to set macros used in a callback. +# By default, the MTA sends all macros defined for a callback. +# If some or all of these are unused, the bandwidth can be saved +# by listing the ones that are used. +# @since 1.0.2 +def symlist(func,*syms): + if func.__name__ not in MACRO_CALLBACKS: + raise ValueError('@symlist applied to non-symlist method: '+func.__name__) + if len(syms) > 5: + raise ValueError('@symlist limited to 5 macros by MTA: '+func.__name__) + func._symlist = syms + return func + ## Disabled action exception. # set_flags() can tell the MTA that this application will not use certain # features (such as CHGFROM). This can also be negotiated for each @@ -393,6 +413,11 @@ class Base(object): def negotiate(self,opts): try: self._actions,p,f1,f2 = opts + for func,stage in MACRO_CALLBACKS.items(): + func = getattr(self,func) + syms = getattr(func,'_symlist',None) + if syms is not None: + self.setsymlist(stage,syms) opts[1] = self._protocol = p & ~self.protocol_mask() opts[2] = 0 opts[3] = 0 @@ -443,23 +468,27 @@ class Base(object): # set. The protocol stages are M_CONNECT, M_HELO, M_ENVFROM, M_ENVRCPT, # M_DATA, M_EOM, M_EOH. # - # May only be called from negotiate callback. + # May only be called from negotiate callback. Hence, this is an advanced + # feature. Use the @@symlist function decorator to conviently set + # the macros used by a callback. # @since 0.9.8, previous version was misspelled! # @param stage the protocol stage to set to macro list for, # one of the M_* constants defined in Milter # @param macros space separated and/or lists of strings def setsymlist(self,stage,*macros): if not self._actions & SETSYMLIST: raise DisabledAction("SETSYMLIST") + if len(macros) > 5: + raise ValueError('setsymlist limited to 5 macros by MTA') a = [] for m in macros: try: m = m.encode('utf8') except: pass try: - m = m.split(' ') + m = m.split(b' ') + a += m except: pass - a += m - return self._ctx.setsymlist(stage,' '.join(a)) + return self._ctx.setsymlist(stage,b' '.join(a)) # Milter methods which can only be called from eom callback. diff --git a/Milter/test.py b/Milter/test.py index f0dc58b62c886769e3b94038650e1fdbd11417fc..c20be6354d4616d2357c5d1de0499760cf137b1a 100644 --- a/Milter/test.py +++ b/Milter/test.py @@ -40,6 +40,9 @@ class TestBase(object): self._reply = None ## The rfc822 message object for the current email being fed to the %milter. self._msg = None + ## The protocol stage for macros returned + self._stage = None + ## The macros returned by protocol stage self._symlist = [ None, None, None, None, None, None, None ] def log(self,*msg): @@ -54,8 +57,12 @@ class TestBase(object): self._macros[name] = val def getsymval(self,name): - # FIXME: track stage, and use _symlist - return self._macros.get(name,'') + stage = self._stage + if stage >= 0: + syms = self._symlist[stage] + if syms is not None and name not in syms: + return None + return self._macros.get(name,None) def replacebody(self,chunk): if self._body: @@ -113,7 +120,10 @@ class TestBase(object): self._reply = (rcode,xcode) + msg def setsymlist(self,stage,macros): - if not self._actions & SETSYMLIST: raise DisabledAction("SETSYMLIST") + if not self._actions & SETSYMLIST: + raise DisabledAction("SETSYMLIST") + if self._stage != -1: + raise RuntimeError("setsymlist may only be called from negotiate") # not used yet, but just for grins we save the data a = [] for m in macros: @@ -121,9 +131,13 @@ class TestBase(object): m = m.encode('utf8') except: pass try: - m = m.split(' ') + m = m.split(b' ') except: pass a += m + if len(a) > 5: + raise ValueError('setsymlist limited to 5 macros by MTA') + if self._symlist[stage] is not None: + raise ValueError('setsymlist already called for stage:'+stage) self._symlist[stage] = set(a) ## Feed a file like object to the %milter. Calls envfrom, envrcpt for @@ -144,16 +158,32 @@ class TestBase(object): self._reply = None self._sender = '<%s>'%sender msg = mime.message_from_file(fp) + # envfrom + self._stage = Milter.M_ENVFROM rc = self.envfrom(self._sender) + self._stage = None if rc != Milter.CONTINUE: return rc + # envrcpt for rcpt in (rcpt,) + rcpts: + self._stage = Milter.M_ENVRCPT rc = self.envrcpt('<%s>'%rcpt) + self._stage = None if rc != Milter.CONTINUE: return rc + # data + self._stage = Milter.M_DATA + rc = self.data() + self._stage = None + if rc != Milter.CONTINUE: return rc + # header for h,val in msg.items(): rc = self.header(h,val) if rc != Milter.CONTINUE: return rc + # eoh + self._stage = Milter.M_EOH rc = self.eoh() + self._stage = None if rc != Milter.CONTINUE: return rc + # body header,body = msg.as_bytes().split(b'\n\n',1) bfp = BytesIO(body) while 1: @@ -163,7 +193,9 @@ class TestBase(object): if rc != Milter.CONTINUE: return rc self._msg = msg self._body = BytesIO() + self._stage = Milter.M_EOM rc = self.eom() + self._stage = None if self._bodyreplaced: body = self._body.getvalue() self._body = BytesIO() @@ -189,12 +221,17 @@ class TestBase(object): self._body = None self._bodyreplaced = False opts = [ Milter.CURR_ACTS,~0,0,0 ] + self._stage = -1 rc = self.negotiate(opts) + self._stage = Milter.M_CONNECT rc = super(TestBase,self).connect(host,1,(ip,1234)) if rc != Milter.CONTINUE: + self._stage = None self.close() return rc + self._stage = Milter.M_HELO rc = self.hello(helo) + self._stage = None if rc != Milter.CONTINUE: self.close() return rc diff --git a/sample.py b/sample.py index 6c48296ba2072c013fa846118e4345633aa37411..4829397c95c30e8f6f0a71afdfb729b237130b6f 100644 --- a/sample.py +++ b/sample.py @@ -33,18 +33,24 @@ class sampleMilter(Milter.Milter): self.fp = None self.bodysize = 0 self.id = Milter.uniqueID() + self.user = None # multiple messages can be received on a single connection # envfrom (MAIL FROM in the SMTP protocol) seems to mark the start # of each message. + @Milter.symlist('{auth_authen}') @Milter.noreply def envfrom(self,f,*str): "start of MAIL transaction" - self.log("mail from",f,str) self.fp = BytesIO() self.tempname = None self.mailfrom = f self.bodysize = 0 + self.user = self.getsymval('{auth_authen}') + if self.user: + self.log("user",self.user,"sent mail from",f,str) + else: + self.log("mail from",f,str) return Milter.CONTINUE def envrcpt(self,to,*str): diff --git a/testsample.py b/testsample.py index 6259257be50560aa91ac34b77eb42acfdaad7b70..2df8c0eb6c567235884a0d3b35034e0e4e342fa9 100644 --- a/testsample.py +++ b/testsample.py @@ -13,9 +13,14 @@ class BMSMilterTestCase(unittest.TestCase): def testDefang(self,fname='virus1'): milter = TestMilter() + milter.setsymval('{auth_authen}','batman') + milter.setsymval('{auth_type}','batcomputer') + milter.setsymval('j','mailhost') rc = milter.connect() self.failUnless(rc == Milter.CONTINUE) rc = milter.feedMsg(fname) + self.failUnless(milter.user == 'batman',"getsymval failed") + self.failUnless(milter.auth_type != 'batcomputer',"setsymlist failed") self.failUnless(rc == Milter.ACCEPT) self.failUnless(milter._bodyreplaced,"Message body not replaced") fp = milter._body