Merge pull request #90 from StrikingLoo/master

refactor limits.py

lib/limits.py

@@ -1,12 +1,12 @@
 """
-Connection limitatation.
+Connection limitation.
 
 Number of connections from one IP is limited.
-We have nothing against scripting and automated queries,
-even in opposite, we encourage them, but there are some
+We have nothing against scripting and automated queries.
+Even the opposite, we encourage them. But there are some
 connection limits that even we can't handle.
-Currently the limits set low, but they will be relaxed
-in future.
+Currently the limits are quite restrictive, but they will be relaxed
+in the future.
 
 Usage:
 
@@ -21,6 +21,19 @@ from globals import log
 
 _WHITELIST = ['5.9.243.177']
 
+
+def _time_caps(m, h, d):
+    return {
+            'min':   m,
+            'hour':  h,
+            'day':   d,
+            }
+
+
+
+
+
+
 class Limits(object):
     """
     Queries limitation (by IP).
@@ -32,28 +45,36 @@ class Limits(object):
 
     def __init__(self):
         self.intervals = ['min', 'hour', 'day']
-        self.divisor = {
-            'min':      60,
-            'hour':     3600,
-            'day':      86400,
-            }
+
+        self.divisor = _time_caps(60, 3600, 86400)
+        self.limit = _time_caps(30, 600, 1000)
+        self.last_update = _time_caps(0, 0, 0)
+
         self.counter = {
             'min':      {},
             'hour':     {},
             'day':      {},
             }
-        self.limit = {
-            'min':      30,
-            'hour':     600,
-            'day':      1000,
-            }
-        self.last_update = {
-            'min':      0,
-            'hour':     0,
-            'day':      0,
-            }
+
         self._clear_counters_if_needed()
 
+     
+    def _log_visit(self, interval, ip_address):
+        if ip_address not in self.counter[interval]:
+            self.counter[interval][ip_address] = 0
+        self.counter[interval][ip_address] += 1
+
+	def _limit_exceeded(self, interval, ip_address):
+		visits = self.counter[interval][ip_address]
+		limit = self._get_limit(interval)
+		return  visits > limit
+
+    def _get_limit(self, interval):
+        return self.limit[interval]
+
+	def _report_excessive_visits(self, interval, ip_address):
+		log("%s LIMITED [%s for %s]" % (ip_address, self._get_limit(interval), interval))
+
     def check_ip(self, ip_address):
         """
         Check if `ip_address` is allowed, and if not raise an RuntimeError exception.
@@ -63,13 +84,11 @@ class Limits(object):
             return None
         self._clear_counters_if_needed()
         for interval in self.intervals:
-            if ip_address not in self.counter[interval]:
-                self.counter[interval][ip_address] = 0
-            self.counter[interval][ip_address] += 1
-            if self.limit[interval] <= self.counter[interval][ip_address]:
-                log("%s LIMITED [%s for %s]" % (ip_address, self.limit[interval], interval))
+            self._log_visit(interval, ip_address)
+            if self._limit_exceeded(interval, ip_address):
+                self._report_excessive_visits(interval, ip_address)
                 return ("Not so fast! Number of queries per %s is limited to %s"
-                        % (interval, self.limit[interval]))
+                        % (interval, self._get_limit(interval) ))
         return None
 
     def reset(self):
@@ -82,6 +101,6 @@ class Limits(object):
     def _clear_counters_if_needed(self):
         current_time = int(time.time())
         for interval in self.intervals:
-            if current_time / self.divisor[interval] != self.last_update[interval]:
+            if current_time // self.divisor[interval] != self.last_update[interval]:
                 self.counter[interval] = {}
                 self.last_update[interval] = current_time / self.divisor[interval]