Repository: awdeorio/mailmerge

Issue:
Rate limiting option
Is there interest in a flood limit timer? For example, my ISP has a limit of 50 emails in 5 minutes.  I'd love to be able to tell mailmerge this so i can just leave a process running.  Perhaps something like:

mailmerge ... --pause-count=50 --pause-time=600

--pause-count = Pause after this many emails are sent - default 0
--pause-time = Pause for this long in seconds - default 0

Error condition if either is set to a negative value or if only one is set to a positive value

WDYT?

Comments:

Patch:
diff --git a/mailmerge/__init__.py b/mailmerge/__init__.py
index dc48419..fd69f6d 100644
--- a/mailmerge/__init__.py
+++ b/mailmerge/__init__.py
@@ -7,4 +7,4 @@ Andrew DeOrio <awdeorio@umich.edu>
 
 from .sendmail_client import SendmailClient
 from .template_message import TemplateMessage
-from .exceptions import MailmergeError
+from .exceptions import MailmergeError, MailmergeRateLimitError
diff --git a/mailmerge/__main__.py b/mailmerge/__main__.py
index bfeb745..d9648db 100644
--- a/mailmerge/__main__.py
+++ b/mailmerge/__main__.py
@@ -5,6 +5,7 @@ Andrew DeOrio <awdeorio@umich.edu>
 """
 from __future__ import print_function
 import sys
+import time
 import codecs
 import textwrap
 import click
@@ -112,9 +113,20 @@ def main(sample, dry_run, limit, no_limit, resume,
         template_message = TemplateMessage(template_path)
         csv_database = read_csv_database(database_path)
         sendmail_client = SendmailClient(config_path, dry_run)
+
         for _, row in enumerate_range(csv_database, start, stop):
             sender, recipients, message = template_message.render(row)
-            sendmail_client.sendmail(sender, recipients, message)
+            while True:
+                try:
+                    sendmail_client.sendmail(sender, recipients, message)
+                except exceptions.MailmergeRateLimitError:
+                    print_bright_white_on_cyan(
+                        ">>> rate limit exceeded, waiting ...",
+                        output_format,
+                    )
+                else:
+                    break
+                time.sleep(1)
             print_bright_white_on_cyan(
                 ">>> message {message_num}"
                 .format(message_num=message_num),
@@ -127,6 +139,7 @@ def main(sample, dry_run, limit, no_limit, resume,
                 output_format,
             )
             message_num += 1
+
     except exceptions.MailmergeError as error:
         hint_text = '\nHint: "--resume {}"'.format(message_num)
         sys.exit(
@@ -217,8 +230,18 @@ def create_sample_input_files(template_path, database_path, config_path):
         """))
     with config_path.open("w") as config_file:
         config_file.write(textwrap.dedent(u"""\
+            # Mailmerge SMTP Server Config
+            # https://github.com/awdeorio/mailmerge
+            #
             # Pro-tip: SSH or VPN into your network first to avoid spam
             # filters and server throttling.
+            #
+            # Parameters
+            #   host       # SMTP server hostname or IP
+            #   port       # SMTP server port
+            #   security   # Security protocol: "SSL/TLS", "STARTTLS", or omit
+            #   username   # Username for SSL/TLS or STARTTLS security
+            #   ratelimit  # Rate limit in messages per minute, 0 for unlimited
 
             # Example: GMail
             [smtp_server]
@@ -226,6 +249,7 @@ def create_sample_input_files(template_path, database_path, config_path):
             port = 465
             security = SSL/TLS
             username = YOUR_USERNAME_HERE
+            ratelimit = 0
 
             # Example: SSL/TLS
             # [smtp_server]
@@ -233,6 +257,7 @@ def create_sample_input_files(template_path, database_path, config_path):
             # port = 465
             # security = SSL/TLS
             # username = YOUR_USERNAME_HERE
+            # ratelimit = 0
 
             # Example: STARTTLS security
             # [smtp_server]
@@ -240,11 +265,13 @@ def create_sample_input_files(template_path, database_path, config_path):
             # port = 25
             # security = STARTTLS
             # username = YOUR_USERNAME_HERE
+            # ratelimit = 0
 
             # Example: No security
             # [smtp_server]
             # host = newman.eecs.umich.edu
             # port = 25
+            # ratelimit = 0
         """))
     print(textwrap.dedent(u"""\
         Created sample template email "{template_path}"
diff --git a/mailmerge/exceptions.py b/mailmerge/exceptions.py
index 7ba05b8..1d29adb 100644
--- a/mailmerge/exceptions.py
+++ b/mailmerge/exceptions.py
@@ -3,3 +3,7 @@
 
 class MailmergeError(Exception):
     """Top level exception raised by mailmerge functions."""
+
+
+class MailmergeRateLimitError(MailmergeError):
+    """Reuse to send message because rate limit exceeded."""
diff --git a/mailmerge/sendmail_client.py b/mailmerge/sendmail_client.py
index cd400fd..129f394 100644
--- a/mailmerge/sendmail_client.py
+++ b/mailmerge/sendmail_client.py
@@ -3,20 +3,26 @@ SMTP client reads configuration and sends message.
 
 Andrew DeOrio <awdeorio@umich.edu>
 """
+import collections
 import socket
 import smtplib
 import configparser
 import getpass
+import datetime
 from . import exceptions
 from . import utils
 
 
+# Type to store info read from config file
+MailmergeConfig = collections.namedtuple(
+    "MailmergeConfig",
+    ["username", "host", "port", "security", "ratelimit"],
+)
+
+
 class SendmailClient(object):
     """Represent a client connection to an SMTP server."""
 
-    # This class is pretty simple.  We don't need more than one public method.
-    # pylint: disable=too-few-public-methods
-    #
     # We need to inherit from object for Python 2 compantibility
     # https://python-future.org/compatible_idioms.html#custom-class-behaviour
     # pylint: disable=bad-option-value,useless-object-inheritance
@@ -24,34 +30,50 @@ class SendmailClient(object):
     def __init__(self, config_path, dry_run=False):
         """Read configuration from server configuration file."""
         self.config_path = config_path
-        self.dry_run = dry_run
-        self.username = None
-        self.password = None
+        self.dry_run = dry_run  # Do not send real messages
+        self.config = None      # Config read from config_path by read_config()
+        self.password = None    # Password read from stdin
+        self.lastsent = None    # Timestamp of last successful send
+        self.read_config()
+
+    def read_config(self):
+        """Read configuration file and return a MailmergeConfig object."""
         try:
-            config = configparser.RawConfigParser()
-            config.read(str(config_path))
-            self.host = config.get("smtp_server", "host")
-            self.port = config.getint("smtp_server", "port")
-            self.security = config.get("smtp_server", "security",
-                                       fallback=None)
-            if self.security == "Never":
-                # Coerce legacy option "security = Never"
-                self.security = None
-            if self.security is not None:
-                # Read username only if needed
-                self.username = config.get("smtp_server", "username")
-        except configparser.Error as err:
+            parser = configparser.RawConfigParser()
+            parser.read(str(self.config_path))
+            host = parser.get("smtp_server", "host")
+            port = parser.getint("smtp_server", "port")
+            security = parser.get("smtp_server", "security", fallback=None)
+            username = parser.get("smtp_server", "username", fallback=None)
+            ratelimit = parser.getint("smtp_server", "ratelimit", fallback=0)
+        except (configparser.Error, ValueError) as err:
             raise exceptions.MailmergeError(
                 "{}: {}".format(self.config_path, err)
             )
 
+        # Coerce legacy option "security = Never"
+        if security == "Never":
+            security = None
+
         # Verify security type
-        if self.security not in [None, "SSL/TLS", "STARTTLS"]:
+        if security not in [None, "SSL/TLS", "STARTTLS"]:
             raise exceptions.MailmergeError(
                 "{}: unrecognized security type: '{}'"
-                .format(self.config_path, self.security)
+                .format(self.config_path, security)
+            )
+
+        # Verify username
+        if security is not None and username is None:
+            raise exceptions.MailmergeError(
+                "{}: username is required for security type '{}'"
+                .format(self.config_path, security)
             )
 
+        # Save validated configuration
+        self.config = MailmergeConfig(
+            username, host, port, security, ratelimit,
+        )
+
     def sendmail(self, sender, recipients, message):
         """Send email message.
 
@@ -62,41 +84,52 @@ class SendmailClient(object):
         if self.dry_run:
             return
 
+        # Check if we've hit the rate limit
+        now = datetime.datetime.now()
+        if self.config.ratelimit and self.lastsent:
+            waittime = datetime.timedelta(minutes=1.0 / self.config.ratelimit)
+            if now - self.lastsent < waittime:
+                raise exceptions.MailmergeRateLimitError()
+
         # Ask for password if necessary
-        if self.security is not None and self.password is None:
+        if self.config.security is not None and self.password is None:
             prompt = ">>> password for {} on {}: ".format(
-                self.username, self.host)
+                self.config.username, self.config.host)
             self.password = getpass.getpass(prompt)
 
         # Send
         try:
             message_flattened = utils.flatten_message(message)
-            if self.security == "SSL/TLS":
-                with smtplib.SMTP_SSL(self.host, self.port) as smtp:
-                    smtp.login(self.username, self.password)
+            host, port = self.config.host, self.config.port
+            if self.config.security == "SSL/TLS":
+                with smtplib.SMTP_SSL(host, port) as smtp:
+                    smtp.login(self.config.username, self.password)
                     smtp.sendmail(sender, recipients, message_flattened)
-            elif self.security == "STARTTLS":
-                with smtplib.SMTP(self.host, self.port) as smtp:
+            elif self.config.security == "STARTTLS":
+                with smtplib.SMTP(host, port) as smtp:
                     smtp.ehlo()
                     smtp.starttls()
                     smtp.ehlo()
-                    smtp.login(self.username, self.password)
+                    smtp.login(self.config.username, self.password)
                     smtp.sendmail(sender, recipients, message_flattened)
-            elif self.security is None:
-                with smtplib.SMTP(self.host, self.port) as smtp:
+            elif self.config.security is None:
+                with smtplib.SMTP(host, port) as smtp:
                     smtp.sendmail(sender, recipients, message_flattened)
         except smtplib.SMTPAuthenticationError as err:
             raise exceptions.MailmergeError(
                 "{}:{} failed to authenticate user '{}': {}"
-                .format(self.host, self.port, self.username, err)
+                .format(host, port, self.config.username, err)
             )
         except smtplib.SMTPException as err:
             raise exceptions.MailmergeError(
                 "{}:{} failed to send message: {}"
-                .format(self.host, self.port, err)
+                .format(host, port, err)
             )
         except socket.error as err:
             raise exceptions.MailmergeError(
                 "{}:{} failed to connect to server: {}"
-                .format(self.host, self.port, err)
+                .format(host, port, err)
             )
+
+        # Update timestamp of last sent message
+        self.lastsent = now
diff --git a/setup.py b/setup.py
index 49636bc..236d946 100644
--- a/setup.py
+++ b/setup.py
@@ -27,6 +27,9 @@ setuptools.setup(
         "click",
         "configparser;python_version<'3.6'",
 
+        # We mock the time when testing the rate limit feature
+        "freezegun",
+
         # The attachments feature relies on a bug fix in the future library
         # https://github.com/awdeorio/mailmerge/pull/56
         "future>0.18.0",


Test patch:
diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py
new file mode 100644
index 0000000..76b29d8
--- /dev/null
+++ b/tests/test_ratelimit.py
@@ -0,0 +1,140 @@
+# coding=utf-8
+# Python 2 source containing unicode https://www.python.org/dev/peps/pep-0263/
+"""
+Tests for SMTP server rate limit feature.
+
+Andrew DeOrio <awdeorio@umich.edu>
+"""
+import textwrap
+import datetime
+import future.backports.email as email
+import future.backports.email.parser  # pylint: disable=unused-import
+import freezegun
+import pytest
+import click
+import click.testing
+from mailmerge import SendmailClient, MailmergeRateLimitError
+from mailmerge.__main__ import main
+
+try:
+    from unittest import mock  # Python 3
+except ImportError:
+    import mock  # Python 2
+
+# Python 2 pathlib support requires backport
+try:
+    from pathlib2 import Path
+except ImportError:
+    from pathlib import Path
+
+# The sh library triggers lot of false no-member errors
+# pylint: disable=no-member
+
+# We're going to use mock_SMTP because it mimics the real SMTP library
+# pylint: disable=invalid-name
+
+
+@mock.patch('smtplib.SMTP')
+def test_sendmail_ratelimit(mock_SMTP, tmp_path):
+    """Verify SMTP library calls."""
+    config_path = tmp_path/"server.conf"
+    config_path.write_text(textwrap.dedent(u"""\
+        [smtp_server]
+        host = open-smtp.example.com
+        port = 25
+        ratelimit = 60
+    """))
+    sendmail_client = SendmailClient(
+        config_path,
+        dry_run=False,
+    )
+    message = email.message_from_string(u"""
+        TO: to@test.com
+        SUBJECT: Testing mailmerge
+        FROM: from@test.com
+
+        Hello world
+    """)
+
+    # First message
+    sendmail_client.sendmail(
+        sender="from@test.com",
+        recipients=["to@test.com"],
+        message=message,
+    )
+    smtp = mock_SMTP.return_value.__enter__.return_value
+    assert smtp.sendmail.call_count == 1
+
+    # Second message exceeds the rate limit, doesn't try to send a message
+    with pytest.raises(MailmergeRateLimitError):
+        sendmail_client.sendmail(
+            sender="from@test.com",
+            recipients=["to@test.com"],
+            message=message,
+        )
+    assert smtp.sendmail.call_count == 1
+
+    # Retry the second message after 1 s because the rate limit is 60 messages
+    # per minute
+    #
+    # Mock the time to be 1.1 s in the future
+    # Ref: https://github.com/spulec/freezegun
+    now = datetime.datetime.now()
+    with freezegun.freeze_time(now + datetime.timedelta(seconds=1)):
+        sendmail_client.sendmail(
+            sender="from@test.com",
+            recipients=["to@test.com"],
+            message=message,
+        )
+    assert smtp.sendmail.call_count == 2
+
+
+@mock.patch('smtplib.SMTP')
+def test_stdout_ratelimit(mock_SMTP, tmpdir):
+    """Verify SMTP server ratelimit parameter."""
+    # Simple template
+    template_path = Path(tmpdir/"mailmerge_template.txt")
+    template_path.write_text(textwrap.dedent(u"""\
+        TO: {{email}}
+        FROM: from@test.com
+
+        Hello world
+    """))
+
+    # Simple database with two entries
+    database_path = Path(tmpdir/"mailmerge_database.csv")
+    database_path.write_text(textwrap.dedent(u"""\
+        email
+        one@test.com
+        two@test.com
+    """))
+
+    # Simple unsecure server config
+    config_path = Path(tmpdir/"mailmerge_server.conf")
+    config_path.write_text(textwrap.dedent(u"""\
+        [smtp_server]
+        host = open-smtp.example.com
+        port = 25
+        ratelimit = 60
+    """))
+
+    # Run mailmerge
+    before = datetime.datetime.now()
+    with tmpdir.as_cwd():
+        runner = click.testing.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            main, [
+                "--no-limit",
+                "--no-dry-run",
+                "--output-format", "text",
+            ]
+        )
+    after = datetime.datetime.now()
+    assert after - before > datetime.timedelta(seconds=1)
+    smtp = mock_SMTP.return_value.__enter__.return_value
+    assert smtp.sendmail.call_count == 2
+    assert result.exit_code == 0
+    # assert result.stderr == ""  # replace when we drop Python 3.4 support
+    assert ">>> message 1 sent" in result.stdout
+    assert ">>> rate limit exceeded, waiting ..." in result.stdout
+    assert ">>> message 2 sent" in result.stdout