Add rate limiting to SMS

docs/api/modules/sms.rst

@@ -4,6 +4,7 @@
 This module provides ".text-<name> <message>" commands that send SMS messages
 via the Twilio api.
 
+
 Config
 ------
 
@@ -19,6 +20,11 @@ Config
             "guy1": "+11111111111",
             "guy2": "+12222222222"
         },
+        "limit": {
+            "enable": true,
+            "period": 900,
+            "max": 5
+        }
     }
 
 .. cmdoption:: account_sid
@@ -45,6 +51,19 @@ Config
 
     Dict of names to phone numbers. Names must be a-zA-Z0-9 and numbers match the format shown above.
 
+.. cmdoption:: limit.enable
+
+    Enable or disable rate limiting. Rate limiting is controlled as a "burst bucket." If enabled, sending an SMS
+    requires 1 and removes one point from the bucket.
+
+.. cmdoption:: limit.period
+
+    Every `limit.period`, a virtual point is added to the bucket.
+
+.. cmdoption:: limit.max
+
+    When adding a point the bucket, the point will be discarded if the bucket already has `limit.max` points.
+
 
 Twilio Setup
 ------------
@@ -56,6 +75,7 @@ URL.
 The webhook listener listens on `/app/gotsms`, so an example webhook URL would
 be `http://1.2.3.4:3000/app/gotsms`.
 
+
 Class Reference
 ---------------
 

examples/data/config/SMS.json

@@ -7,5 +7,10 @@
     "contacts": {
         "guy1": "+11111111111",
         "guy2": "+12222222222"
+    },
+    "limit": {
+        "enable": true,
+        "period": 900,
+        "max": 5
     }
 }

pyircbot/modules/SMS.py

@@ -5,6 +5,8 @@
     :synopsis: SMS client script (requires Twilio account)
 """
 
+from time import time
+from math import floor
 from pyircbot.modulebase import ModuleBase, regex
 import cherrypy
 from threading import Thread
@@ -62,6 +64,39 @@ class SMS(ModuleBase):
         self.apithread = None
         self.twilio = Client(self.config["account_sid"], self.config["auth_token"])
 
+        # limit-related vars
+        # How many messages can be bursted
+        self.bucket_max = int(self.config["limit"]["max"])
+        # burst bucket, initial value is 1 or half the max, whichever is more
+        self.bucket = max(1, self.bucket_max / 2)
+        # how often the bucket has 1 item added
+        self.bucket_period = int(self.config["limit"]["period"])
+        # last time the burst bucket was filled
+        self.bucket_lastfill = int(time())
+
+    def check_rate_limit(self):
+        """
+        Rate limiting via a 'burst bucket'. This method is called before sending and returns true or false depending on
+        if the action is allowed.
+        """
+
+        # First, update the bucket
+        # Check if $period time has passed since the bucket was filled
+        since_fill = int(time()) - self.bucket_lastfill
+        if since_fill > self.bucket_period:
+            # How many complete points are credited
+            fills = floor(since_fill / self.bucket_period)
+            self.bucket += fills
+            if self.bucket > self.bucket_max:
+                self.bucket = self.bucket_max
+            # Advance the lastfill time appropriately
+            self.bucket_lastfill += self.bucket_period * fills
+
+        if self.bucket >= 1:
+            self.bucket -= 1
+            return True
+        return False
+
     def api(self):
         """
         Run the webhook listener and block
@@ -113,6 +148,11 @@ class SMS(ModuleBase):
         if contact not in self.config["contacts"].keys():
             return  # TODO invalid contact
 
+        if self.config["limit"]["enable"]:
+            if not self.check_rate_limit():
+                self.bot.act_PRIVMSG(msg.args[0], "Sorry, try again later")
+                return
+
         try:
             self.twilio.api.account.messages.create(to=self.config["contacts"][contact],
                                                     from_=self.config["number"],