def tell(self, msg, action, learnas=''): """Tell what type of we are to process and what should be done with that message. This includes setting or removing a local or a remote database (learning, reporting, forgetting, revoking).""" action = _check_action(action) mode = learnas.upper() headers = { 'Message-class': '', 'Set': 'local', } if action == 'learn': if mode == 'SPAM': headers['Message-class'] = 'spam' elif mode in ['HAM', 'NOTSPAM', 'NOT_SPAM']: headers['Message-class'] = 'ham' else: raise SpamCError('The learnas option is invalid') elif action == 'forget': del headers['Message-class'] del headers['Set'] headers['Remove'] = 'local' elif action == 'report': headers['Message-class'] = 'spam' headers['Set'] = 'local, remote' elif action == 'revoke': headers['Message-class'] = 'ham' headers['Remove'] = 'remote' return self.perform('TELL', msg, headers)
def _check_action(action): """check for invalid actions""" if isinstance(action, types.StringTypes): action = action.lower() if action not in ['learn', 'forget', 'report', 'revoke']: raise SpamCError('The action option is invalid') return action
def learn(self, msg, learnas): """Learn message as spam/ham or forget""" if not isinstance(learnas, types.StringTypes): raise SpamCError('The learnas option is invalid') if learnas.lower() == 'forget': resp = self.tell(msg, 'forget') else: resp = self.tell(msg, 'learn', learnas) return resp
def perform(self, cmd, msg='', extra_headers=None): """Perform the call""" tries = 0 while 1: conn = None try: conn = self.get_connection() if hasattr(msg, 'read') and hasattr(msg, 'fileno'): msg_length = str(os.fstat(msg.fileno()).st_size) elif hasattr(msg, 'read'): msg.seek(0, 2) msg_length = str(msg.tell() + 2) else: if msg: try: msg_length = str(len(msg) + 2) except TypeError: conn.close() raise ValueError( 'msg param should be a string or file handle') else: msg_length = '2' headers = self.get_headers(cmd, msg_length, extra_headers) if isinstance(msg, types.StringTypes): if self.gzip and msg: msg = compress(msg + '\r\n', self.compress_level) else: msg = msg + '\r\n' conn.send(headers + msg) else: conn.send(headers) if hasattr(msg, 'read'): if hasattr(msg, 'seek'): msg.seek(0) conn.sendfile(msg, self.gzip, self.compress_level) conn.send('\r\n') try: conn.socket().shutdown(socket.SHUT_WR) except socket.error: pass return get_response(cmd, conn) except socket.gaierror, err: if conn is not None: conn.release() raise SpamCError(str(err)) except socket.timeout, err: if conn is not None: conn.release() raise SpamCTimeOutError(str(err))
class SpamC(object): """Spamc Client class""" # pylint: disable=R0913 def __init__(self, host=None, port=783, socket_file='/var/run/spamassassin/spamd.sock', user=None, timeout=None, wait_tries=0.3, max_tries=3, backend="thread", gzip=None, compress_level=6, is_ssl=None, **ssl_args): """Init""" self.host = host self.port = port self.socket_file = socket_file self.user = user if isinstance(backend, str): self.backend_mod = load_backend(backend) else: self.backend_mod = backend self.max_tries = max_tries self.wait_tries = wait_tries self.timeout = timeout self.gzip = gzip self.compress_level = compress_level self.is_ssl = is_ssl self.ssl_args = ssl_args or {} def get_connection(self): """Creates a new connection""" if self.host is None: connector = SpamCUnixConnector conn = connector(self.socket_file, self.backend_mod) else: connector = SpamCTcpConnector conn = connector(self.host, self.port, self.backend_mod, is_ssl=self.is_ssl, **self.ssl_args) return conn def get_headers(self, cmd, msg_length, extra_headers): """Returns the headers string based on command to execute""" cmd_header = "%s %s" % (cmd, PROTOCOL_VERSION) len_header = "Content-length: %s" % msg_length headers = [cmd_header, len_header] if self.user: user_header = "User: %s" % self.user headers.append(user_header) if self.gzip: headers.append("Compress: zlib") if extra_headers is not None: for key in extra_headers: if key.lower() != 'content-length': headers.append("%s: %s" % (key, extra_headers[key])) headers.append('') headers.append('') return '\r\n'.join(headers) # pylint: disable=E1103 def perform(self, cmd, msg='', extra_headers=None): """Perform the call""" tries = 0 while 1: conn = None try: conn = self.get_connection() if hasattr(msg, 'read') and hasattr(msg, 'fileno'): msg_length = str(os.fstat(msg.fileno()).st_size) elif hasattr(msg, 'read'): msg.seek(0, 2) msg_length = str(msg.tell() + 2) else: if msg: try: msg_length = str(len(msg) + 2) except TypeError: conn.close() raise ValueError( 'msg param should be a string or file handle') else: msg_length = '2' headers = self.get_headers(cmd, msg_length, extra_headers) if isinstance(msg, types.StringTypes): if self.gzip and msg: msg = compress(msg + '\r\n', self.compress_level) else: msg = msg + '\r\n' conn.send(headers + msg) else: conn.send(headers) if hasattr(msg, 'read'): if hasattr(msg, 'seek'): msg.seek(0) conn.sendfile(msg, self.gzip, self.compress_level) conn.send('\r\n') try: conn.socket().shutdown(socket.SHUT_WR) except socket.error: pass return get_response(cmd, conn) except socket.gaierror, err: if conn is not None: conn.release() raise SpamCError(str(err)) except socket.timeout, err: if conn is not None: conn.release() raise SpamCTimeOutError(str(err)) except socket.error, err: if conn is not None: conn.close() errors = (errno.EAGAIN, errno.EPIPE, errno.EBADF, errno.ECONNRESET) if err[0] not in errors or tries >= self.max_tries: raise SpamCError("socket.error: %s" % str(err))