1. George Notaras
  2. simsmtp

Source

simsmtp / src / simsmtp / main.py

# -*- coding: utf-8 -*-
#
#  This file is part of simsmtp.
#
#  simsmtp - 
#
#  Project: https://www.codetrax.org/projects/simsmtp
#
#  Copyright 2009 George Notaras <gnot [at] g-loaded.eu>, CodeTRAX.org
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#


import sys
import os
import smtplib
import logging
import socket
import threading
import time
import random
from optparse import OptionParser

import mimetypes
from email.MIMEMultipart import MIMEMultipart
from email.MIMEBase import MIMEBase
from email.MIMEText import MIMEText
from email import Encoders
from email.utils import make_msgid

import info

__version__ = info.version
__program__ = info.name
#__program__ = os.path.basename(sys.argv[0])


class SMTPClient(threading.Thread):
    
    def __init__(self, host='localhost', port=25, sender_pool=[], recipient_pool=[], msg_body_pool=[], attachment_pool=[], nrmsg=1):
        threading.Thread.__init__(self)
        self.host = host
        self.port = port
        self.sender_pool = sender_pool
        self.recipient_pool = recipient_pool
        self.msg_body_pool = msg_body_pool
        self.attachment_pool = attachment_pool
        self.nrmsg = nrmsg    # number of messages to send
        
        # Statistics
        self.successful = 0        # counter
        self.longest_send = 0    # gauge
    
    def run(self):
        # If more than one thread are used and if delay has been specified with the "min,max"
        # format (which implies randomness in the delay), delay here, so that not all clients
        # start sending at the same moment. If a static delay has been set, then there is no point
        # in delaying here, as all clients will start sending simultaneusly anyway.
        if options.count > 1 and len(options.delay) > 1:
            delay(self.name)
            
        for n in range(self.nrmsg):
            try:
                t_start = time.time()
                self.sendmail()
                t_end = time.time()
            except socket.error, e:
                logger.critical( '%s ABORT ALL SENDS! %s' % (self.name, e) )
                return
            except smtplib.SMTPException, e:
                logger.error( '%s failed to send email #%d %s' % (self.name, n+1, e) )
            else:
                self.set_stats(t_start, t_end, success=True)
                logger.info( '%s sent email #%d of %d' % (self.name, n+1, self.nrmsg) )
            if n+1 != self.nrmsg:    # do not delay after sending the last message
                delay(self.name)
        logger.info( '%s finished' % self.name )
    
    def set_stats(self, t_start, t_end, success=True):
        if success:
            self.successful += 1
            t_send = t_end - t_start
            if t_send > self.longest_send:
                self.longest_send = t_send
            
    def sendmail(self):
        from_addr = random.choice(self.sender_pool)
        to_addr = random.choice(self.recipient_pool)
        # Ensure from_addr != to_addr
        #if len(self.recipient_pool) > 1:    # TODO: this is probably incomplete. An endless loop might be hiding in here...
        #    # Loop until recipient is different from sender in the case the same file is used for senders and recipients
        #    while from_addr == to_addr:
        #        to_addr = random.choice(self.recipient_pool)
        msg_txt = self.get_email_message(from_addr, to_addr)
        server = smtplib.SMTP(self.host, self.port)
        if options.localhostname:
            server.helo(options.localhostname)
        server.sendmail(from_addr, to_addr, msg_txt)
        server.quit()
    
    def get_email_message(self, from_addr, to_addr):
        """
        Returns the message object as string
        """
        if not options.subject:
            options.subject = 'Stress Test'
        subject = '[%s][%s] -- %s' % (PID, self.name, options.subject)
        
        if self.attachment_pool:
            # Create the container email message.
            msg_obj = MIMEMultipart()
            msg_obj["From"] = from_addr
            msg_obj["To"] = to_addr
            msg_obj["Subject"] = subject
            msg_obj["Message-Id"] = make_msgid()
            msg_obj.preamble = ''
            # Guarantees the message ends in a newline
            msg_obj.epilogue = ''
            
            # TEXT PART
            body, filename = self.get_content_from_random_file(self.msg_body_pool)
            msg_text = MIMEText(body, "plain", "utf-8")
            msg_obj.attach(msg_text)
            
            # ATTACHMENT PART
            file_data, filename = self.get_content_from_random_file(self.attachment_pool.keys())
            mime_type, mime_subtype = self.attachment_pool[filename]
            msg_atch = MIMEBase(mime_type, mime_subtype)
            msg_atch.set_payload(file_data)
            # Encode the payload using Base64
            Encoders.encode_base64(msg_atch)
            # Set the filename parameter
            msg_atch.add_header('Content-Disposition', 'attachment', filename=os.path.basename(filename))
            msg_obj.attach(msg_atch)
        else:
            # Simple text message
            body, filename = self.get_content_from_random_file(self.msg_body_pool)
            msg_obj = MIMEText(body, 'plain', 'utf-8')
            msg_obj['From'] = from_addr
            msg_obj['To'] = to_addr
            msg_obj['Subject'] = subject
            msg_obj["Message-Id"] = make_msgid()
        return msg_obj.as_string()
    
    def get_content_from_random_file(self, filename_pool):
        contents = ''
        if filename_pool:
            fpath = random.choice(filename_pool)
            f = open(fpath, 'rb')
            contents = f.read()
            f.close()
        return contents, fpath


def parseargs():
    parser = OptionParser()
    parser.add_option('-f', '--sender', dest='sender', help='The sender\'s email address. A path to a file containing email addresses (one per line) can be used. In that case, a random sender will be chosen for each email.', metavar='ADDRESS or PATH')
    parser.add_option('-r', '--recipient', dest='recipient', help='The recipient\'s email address. A path to a file containing email addresses (one per line) can be used. In that case, a random recipient will be chosen for each email.', metavar='ADDRESS or PATH')
    parser.add_option('-m', '--message', dest='message', help='Path to a text file which will be used as the message body. If the provided path is a directory, a random choice of the files inside the directory will be made. No check will be performed whether the file is a text file or not!!!', metavar='PATH')
    parser.add_option('-a', '--attachment', dest='attachment', default='', help='Path to a file which will be attached to the emails. If the provided path is a directory, a random choice of the files inside the directory will be made and one of them will be finally attached to the email. No kinds of checks will be performed on the file(s)!!!', metavar='PATH')
    parser.add_option('-s', '--subject', dest='subject', help='A subject that will be set to the email messages. Enclose in quotes if it contains spaces -- very likely.', metavar='"SUBJECT TEXT"')
    parser.add_option('-k', '--host', dest='host', default='localhost', help='The server\'s hostname or IP address.', metavar='HOSTNAME')
    parser.add_option('-p', '--port', dest='port', type="int", default=25, help='The SMTP server\'s port', metavar='PORT')
    parser.add_option('-c', '--count', dest='count', type="int", default=1, help='The number of threads to start.', metavar='NUMBER')
    parser.add_option('-n', '--number-of-msgs', dest='nrmsg', type="int", default=1, help='The number of messages each thread will send.', metavar='NUMBER')
    parser.add_option('-d', '--delay', dest='delay', default='1', help='The number of seconds to wait between sends. The delay can be specified as a "min,max" delay (no spaces). In that case, a random delay between min & max will be used.', metavar='SECONDS')
    parser.add_option('-l', '--local-hostname', dest='localhostname', default='localhost', help='If specified, this hostname will be sent to the remote server in the HELO command.', metavar='HOSTNAME')
    options, args = parser.parse_args()
    if args:
        parser.error('Wrong number of arguments')
    elif not options.sender:
        parser.error('Sender is required')
    elif not options.recipient:
        parser.error('Recipient is required')
    elif not options.message:
        parser.error('Path to message or path to dir with messages is required')
    elif options.delay:
        # Check if specified delays are integers and create a list
        try:
            options.delay = [int(i) for i in options.delay.split(',')]
        except ValueError:
            parser.error('Invalid delay')
        else:
            if len(options.delay) not in (1, 2):
                parser.error('Invalid delay')
    return options


def get_logger():
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_formatter = logging.Formatter('%(asctime)s -- %(message)s', '%Y-%m-%d %H:%M:%S')
    stderr_handler.setFormatter(stderr_formatter)
    logger.addHandler(stderr_handler)
    return logger


def get_pool_from_file(option):
    path = os.path.abspath(option)
    file_pool = []
    if os.path.isfile(path):
        f = open(path, 'rb')
        for line in f:
            line = line.strip()
            if line:
                file_pool.append(line)
        f.close()
    else:
        file_pool.append(option)
    return file_pool


def get_msg_body_pool():
    path = os.path.abspath(options.message)
    msg_body_pool = []
    if os.path.isfile(path):
        msg_body_pool.append(path)
    elif os.path.isdir(path):
        for fname in os.listdir(path):
            fname_path = os.path.join(path, fname)
            if os.path.isfile(fname_path):
                msg_body_pool.append(fname_path)
    return msg_body_pool


def get_attachment_pool():
    """
    Attachment pool is a dictionary. Format:
    { '<path_to_file>' : <mimetype as tuple (type, subtype)>, }
    
    Example:
    { '/tmp/readme.txt' : ('text', 'plain'), }
    """
    attachment_pool = {}
    if not options.attachment:
        return attachment_pool
    
    def add_to_pool(path):
        mime_type, encoding = mimetypes.guess_type(path)
        attachment_pool[path] = mime_type.split('/')
    
    path = os.path.abspath(options.attachment)
    if os.path.isfile(path):
        add_to_pool(path)
    elif os.path.isdir(path):
        for fname in os.listdir(path):
            fname_path = os.path.join(path, fname)
            if os.path.isfile(fname_path):
                add_to_pool(fname_path)
    return attachment_pool


def delay(name):
    # options.delay has been transformed to a list in parseargs()
    if len(options.delay) > 1:
        n = random.randint(options.delay[0], options.delay[-1])
    else:
        n = options.delay[0]
    logger.debug('%s sleeping for %s second(s)' % (name, n))
    time.sleep(n)


def stress():
    pool = {}    # smtp client pool
    sender_pool = get_pool_from_file(options.sender)
    recipient_pool = get_pool_from_file(options.recipient)
    msg_body_pool = get_msg_body_pool()
    attachment_pool = get_attachment_pool()
    # print sender_pool
    # print recipient_pool
    # print msg_body_pool
    # print attachment_pool
    for n in range(options.count):
        name = 'smtp-%s' % str(n+1).zfill(len(str(options.count)))
        pool[name] = SMTPClient(
            host = options.host,
            port = options.port,
            sender_pool = sender_pool,
            recipient_pool = recipient_pool,
            msg_body_pool = msg_body_pool,
            attachment_pool = attachment_pool,
            nrmsg=options.nrmsg
        )
        pool[name].name = name
        pool[name].start()

    for k in pool.keys():
        pool[k].join()
    
    summary(pool)


def summary(pool):
    print
    print 'Results'
    print '-'*40
    tot_success = 0
    longest_send = 0
    for k in pool.keys():
        tot_success += pool[k].successful
        if pool[k].longest_send > longest_send:
            longest_send = pool[k].longest_send
        print '%s %s/%s' % (pool[k].name.ljust(10), pool[k].successful, options.nrmsg)
    print '-'*40
    print 'Total success: %s/%s' % (tot_success, options.count * options.nrmsg)
    print 'Longest send : %.3f sec' % longest_send


def main():
    t_start = time.time()
    print
    print 'Using %s v%s' % (__program__, __version__)
    print 'Attempting a total of %d messages using %d threads' % (options.count * options.nrmsg, options.count)
    print
    stress()
    t_finish = time.time()
    print 'Elapsed %.3f seconds' % (t_finish - t_start)


options = parseargs()
logger = get_logger()
PID = os.getpid()    # for use in the subject

if __name__ == '__main__':
    main()