milter-sssp

#!/usr/bin/python
#
# Milter to scan mail via SSSP for viruses
#
#   Copyright 2018,2019 Andreas Thienemann <andreas@bawue.net>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

import argparse
import email
import os
import socket
import sys
from multiprocessing import Process
import Milter
from Milter.utils import parse_addr
import syslog
import sssp
import textwrap
import traceback

milter_name = 'milter-sssp'
milter_version = '0.1.0'

global args

class ssspMilter(Milter.Base):
  def __init__(self):  # A new instance with each new connection.
    self.queueid = None
    self.id = Milter.uniqueID()  # Integer incremented with each call.
    self.rcpts = []

  def log(self, msg):
    if self.queue_id:
      syslog.syslog('{}: {}'.format(self.queue_id, msg))
    else:
      syslog.syslog('{}: {}'.format(self.id, msg))

  # multiple messages can be received on a single connection
  # envfrom (MAIL FROM in the SMTP protocol) seems to mark the start
  # of each message.
  def envfrom(self, envfrom, *str):
    self.mail = ''
    if envfrom.startswith('<') and envfrom.endswith('>'):
      self.env_from = envfrom[1:-1]
    else:
      self.env_from = envfrom
    return Milter.CONTINUE

  def envrcpt(self, to, *str):
    if to.startswith('<') and to.endswith('>'):
      self.rcpts.append(to[1:-1])
    else:
      self.rcpts.append(to)
    return Milter.CONTINUE

  def header(self, name, val):
    self.mail += '{}: {}\n'.format(name, val)
    return Milter.CONTINUE

  def eoh(self):
    self.mail += '\n'
    return Milter.CONTINUE

  def body(self, chunk):
    self.mail += chunk
    return Milter.CONTINUE

  def eom(self):
    self.queue_id = self.getsymval('i')
    try:
      scanner = sssp.sssp(args.sssp_socket)
      if not scanner.selftest():
        self.log('SAVDI selftest failed. Not scanning, accepting mail.')
        return Milter.ACCEPT
      result, msg = scanner.check(self.mail)
    except:
      self.log('Unknown SAVDI Error. {}'.format(traceback.format_exc()))
      return Milter.ACCEPT

    engine = scanner.query_engine()
    server = scanner.query_server()
    self.addheader('X-Virus-Scanner', textwrap.fill('{}, Engine: {}, SAV: {} ({}) on {} using {} {}'.format(server['version'],
                     engine['engineversion'], engine['savversion'], engine['virusdatachecksum'],
                     socket.gethostname(), milter_name, milter_version), 61), -1)
    if result:
      self.addheader('X-Virus-Scan', 'Found to be clean.', -1)
      self.log('{}, accepting.'.format(msg))
    else:
      self.addheader('X-Virus-Scan', 'Found to be infected. ({})'.format(msg.split()[-1]), -1)
      self.log('Mail reported infected with {} by SAVDI.'.format(msg.split()[-1]))
      if args.quarantine:
        self.log('{}, marked as hold.'.format(msg))
        self.quarantine(msg)
        return Milter.ACCEPT
      if args.reject:
        self.log('{}, rejecting.'.format(msg))
        self.setreply('550', '5.7.1', 'Message rejected due to malware/virus scanning.')
        return Milter.REJECT
    return Milter.ACCEPT

def main():
  global args

  syslog.openlog(ident='milter-sssp', logoption=syslog.LOG_PID, facility=syslog.LOG_MAIL)
  syslog.syslog('Milter starting using socket {}'.format(args.socket))
  job = Process()
  job.start()
  Milter.factory = ssspMilter
  timeout = 60
  Milter.runmilter('ssspMilter', args.socket, timeout)
  job.join()
  syslog.syslog('Milter stopping')
  syslog.closelog()

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='Milter to scan mail for viruses via SSSP.')
  parser.add_argument('-q', '--quarantine', action='store_true', default=False, help='Quarantine suspect mail rather than rejecting it.')
  parser.add_argument('-r', '--reject', action='store_true', default=False, help='Reject suspect mail rather than accepting and marking it.')
  parser.add_argument('socket', help='Milter socket for communicating to postfix')
  parser.add_argument('sssp_socket', help='Socket for communicating to sssp interface.')
  args = parser.parse_args()

  main()