1
1
mirror of https://github.com/chubin/cheat.sh.git synced 2024-12-12 06:24:12 +03:00

Merge pull request #90 from StrikingLoo/master

refactor limits.py
This commit is contained in:
Igor Chubin 2018-08-06 23:59:16 +02:00 committed by GitHub
commit 13c440e0e7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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]