├── tests ├── __init__.py ├── urls.py ├── views.py ├── middleware.py └── test_ratelimiter.py ├── gae_django_ratelimiter ├── views.py ├── models.py ├── __init__.py ├── settings.py └── ratelimiter.py ├── requirements-dev.txt ├── .travis.yml ├── setup_test.sh ├── LICENSE ├── .gitignore └── README.md /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /gae_django_ratelimiter/views.py: -------------------------------------------------------------------------------- 1 | # Intentionally left blank. 2 | -------------------------------------------------------------------------------- /gae_django_ratelimiter/models.py: -------------------------------------------------------------------------------- 1 | # Intentionally left blank. 2 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | flake8 2 | mock 3 | coveralls 4 | pyyaml 5 | -------------------------------------------------------------------------------- /gae_django_ratelimiter/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.1.0' # noqa 2 | 3 | from .ratelimiter import ratelimit, RateLimiterMiddleware, HttpResponseThrottled # noqa 4 | -------------------------------------------------------------------------------- /tests/urls.py: -------------------------------------------------------------------------------- 1 | from django.conf.urls import url 2 | import views 3 | 4 | urlpatterns = [ 5 | url(r'^random/?$', views.random, name='random'), 6 | url(r'^notrandom/?$', views.notrandom, name='notrandom'), 7 | url(r'^decorated/?$', views.decorated, name='decorated'), 8 | ] 9 | -------------------------------------------------------------------------------- /tests/views.py: -------------------------------------------------------------------------------- 1 | from django.http import HttpResponse 2 | from gae_django_ratelimiter import ratelimit 3 | 4 | 5 | def random(request): 6 | return HttpResponse('random') 7 | 8 | 9 | def notrandom(request): 10 | return HttpResponse('notrandom') 11 | 12 | 13 | @ratelimit(prefix='gaerl', minutes=1, requests=5) 14 | def decorated(request): 15 | return HttpResponse('notrandom') 16 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | 3 | python: 4 | - "2.7" 5 | 6 | env: 7 | - DJANGO_VER="1.9" 8 | - DJANGO_VER="1.11" 9 | 10 | install: 11 | - bash setup_test.sh 12 | - pip install django==$DJANGO_VER -q 13 | - pip install -r requirements-dev.txt -q 14 | 15 | script: 16 | - flake8 gae_django_ratelimiter tests 17 | - coverage run --omit=*/__init__.py --source=gae_django_ratelimiter -m unittest discover tests -v 18 | 19 | after_success: 20 | - coveralls 21 | 22 | branches: 23 | only: 24 | - master 25 | -------------------------------------------------------------------------------- /setup_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | 5 | # Source: https://blog.bekt.net/p/gae-sdk/ 6 | API_CHECK='https://appengine.google.com/api/updatecheck' 7 | SDK_VERSION=$(curl --retry 2 -s $API_CHECK | awk -F '\"' '/release/ {print $2}') 8 | # Remove the dots. 9 | SDK_VERSION_S="${SDK_VERSION//./}" 10 | 11 | SDK_URL='https://storage.googleapis.com/appengine-sdks/' 12 | SDK_URL_A="${SDK_URL}featured/google_appengine_${SDK_VERSION}.zip" 13 | SDK_URL_B="${SDK_URL}deprecated/$SDK_VERSION_S/google_appengine_${SDK_VERSION}.zip" 14 | 15 | function download_sdk { 16 | echo ">>> Downloading... GAE ver $SDK_VERSION" 17 | mkdir -p gae_sdk 18 | curl -L --retry 2 -s --fail -o "gae.zip" "$SDK_URL_A" || \ 19 | curl -L --retry 2 --fail -s -o "gae.zip" "$SDK_URL_B" || \ 20 | exit 1 21 | unzip -qd gae_sdk "gae.zip" && rm gae.zip 22 | } 23 | 24 | download_sdk 25 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 ping https://github.com/ping/ 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /tests/middleware.py: -------------------------------------------------------------------------------- 1 | from gae_django_ratelimiter import RateLimiterMiddleware 2 | 3 | 4 | class TestRateLimiterMiddleware(RateLimiterMiddleware): 5 | requests = 3 6 | minutes = 0.5 7 | prefix = 'xyz' 8 | 9 | 10 | class DisabledRateLimiterMiddleware(RateLimiterMiddleware): 11 | enabled = False 12 | 13 | 14 | class CooldownRateLimiterMiddleware(RateLimiterMiddleware): 15 | requests = 2 16 | minutes = 2.0/60 17 | cooldown_from_last_request = True 18 | cooldown_minutes = 0.1 19 | 20 | 21 | class ExcludeRateLimiterMiddleware(RateLimiterMiddleware): 22 | exclude_url_names = ['random'] 23 | 24 | 25 | class IncludeRateLimiterMiddleware(RateLimiterMiddleware): 26 | include_url_names = ['random'] 27 | 28 | 29 | class ExcludeAuthRateLimiterMiddleware(RateLimiterMiddleware): 30 | exclude_authenticated = True 31 | exclude_admins = True 32 | 33 | 34 | class IncludeAuthRateLimiterMiddleware(RateLimiterMiddleware): 35 | exclude_authenticated = False 36 | exclude_admins = False 37 | 38 | 39 | class ExcludeAuthAdminRateLimiterMiddleware(RateLimiterMiddleware): 40 | exclude_authenticated = False 41 | exclude_admins = True 42 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.sh 2 | .vscode/ 3 | gae_sdk/ 4 | venv/ 5 | 6 | *.iml 7 | .idea/ 8 | 9 | # OS generated files # 10 | .DS_Store 11 | .DS_Store? 12 | ._* 13 | .Spotlight-V100 14 | .Trashes 15 | ehthumbs.db 16 | Thumbs.db 17 | 18 | # Byte-compiled / optimized / DLL files 19 | __pycache__/ 20 | *.py[cod] 21 | *$py.class 22 | 23 | # C extensions 24 | *.so 25 | 26 | # Distribution / packaging 27 | .Python 28 | env/ 29 | build/ 30 | develop-eggs/ 31 | dist/ 32 | downloads/ 33 | eggs/ 34 | .eggs/ 35 | lib/ 36 | lib64/ 37 | parts/ 38 | sdist/ 39 | var/ 40 | *.egg-info/ 41 | .installed.cfg 42 | *.egg 43 | 44 | # PyInstaller 45 | # Usually these files are written by a python script from a template 46 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 47 | *.manifest 48 | *.spec 49 | 50 | # Installer logs 51 | pip-log.txt 52 | pip-delete-this-directory.txt 53 | 54 | # Unit test / coverage reports 55 | htmlcov/ 56 | .tox/ 57 | .coverage 58 | .coverage.* 59 | .cache 60 | nosetests.xml 61 | coverage.xml 62 | *,cover 63 | .hypothesis/ 64 | 65 | # Translations 66 | *.mo 67 | *.pot 68 | 69 | # Django stuff: 70 | *.log 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | target/ 77 | 78 | #Ipython Notebook 79 | .ipynb_checkpoints 80 | -------------------------------------------------------------------------------- /gae_django_ratelimiter/settings.py: -------------------------------------------------------------------------------- 1 | from django.conf import settings 2 | 3 | RATELIMITER_ENABLED = getattr( 4 | settings, 'GAE_DJANGO_RATELIMITER_ENABLED', True) 5 | RATELIMITER_CACHE_PREFIX = getattr( 6 | settings, 'GAE_DJANGO_RATELIMITER_CACHE_PREFIX', 'gaerl') 7 | 8 | RATELIMITER_CACHE_MINUTES = getattr( 9 | settings, 'GAE_DJANGO_RATELIMITER_CACHE_MINUTES', 2) 10 | RATELIMITER_CACHE_REQUESTS = getattr( 11 | settings, 'GAE_DJANGO_RATELIMITER_CACHE_REQUESTS', 20) 12 | 13 | RATELIMITER_COOLDOWN_LAST_REQ = getattr( 14 | settings, 'GAE_DJANGO_RATELIMITER_COOLDOWN_LAST_REQ', False) 15 | 16 | RATELIMITER_COOLDOWN_MINUTES = getattr( 17 | settings, 'GAE_DJANGO_RATELIMITER_COOLDOWN_MINUTES', 18 | RATELIMITER_CACHE_MINUTES) 19 | 20 | # Expects a list of url names (urlconf) 21 | RATELIMITER_INCLUDE_URL_NAMES = getattr( 22 | settings, 'GAE_DJANGO_RATELIMITER_INCLUDE_URL_NAMES', []) 23 | # Expects a list of url names (urlconf) 24 | RATELIMITER_EXCLUDE_URL_NAMES = getattr( 25 | settings, 'GAE_DJANGO_RATELIMITER_EXCLUDE_URL_NAMES', []) 26 | 27 | RATELIMITER_EXCLUDE_AUTHENTICATED = getattr( 28 | settings, 'GAE_DJANGO_RATELIMITER_EXCLUDE_AUTHENTICATED', True) 29 | 30 | RATELIMITER_EXCLUDE_ADMINS = getattr( 31 | settings, 'GAE_DJANGO_RATELIMITER_EXCLUDE_ADMINS', True) 32 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # gae-django-ratelimiter 2 | 3 | A basic Django request rate limiter (X requests per Y minutes) for use on Google App Engine (GAE). 4 | 5 | Just Good Enough™ for weekend projects. 6 | 7 | [![Build](https://img.shields.io/travis/ping/gae-django-ratelimiter.svg)](https://travis-ci.org/ping/gae-django-ratelimiter) 8 | [![Coverage](https://img.shields.io/coveralls/ping/gae-django-ratelimiter.svg)](https://coveralls.io/github/ping/gae-django-ratelimiter) 9 | 10 | ## Features 11 | 12 | 1. Available as middleware ``gae_django_ratelimiter.RateLimiterMiddleware`` or request decorator ``@ratelimit`` 13 | 1. Does not limit GAE tasks by default 14 | 1. Does not do much else 15 | 16 | Compatible with Django versions >=1.4, <=1.11 17 | 18 | ## Usage 19 | 20 | ### Basic 21 | 22 | 1. Add the ``gae_django_ratelimiter`` folder to your project path 23 | 24 | 1. Use either the decorator for individual views OR include it in your ``INSTALLED_APPS`` and ``MIDDLEWARE_CLASSES`` (Django<=1.9) or ``MIDDLEWARE`` (Django>=1.10) in your app's ``settings.py`` to apply it globally. 25 | 26 | ```python 27 | # views.py 28 | from gae_django_ratelimiter import ratelimit 29 | 30 | @ratelimit(requests=30, minutes=2) 31 | def a_view(request): 32 | # ... 33 | ``` 34 | 35 | ```python 36 | # settings.py 37 | 38 | # Custom Configuration 39 | GAE_DJANGO_RATELIMITER_ENABLED = True 40 | GAE_DJANGO_RATELIMITER_CACHE_PREFIX = 'lolcat' 41 | GAE_DJANGO_RATELIMITER_CACHE_MINUTES = 2 42 | GAE_DJANGO_RATELIMITER_CACHE_REQUESTS = 30 43 | 44 | # For Django <= 1.9. For newer versions set MIDDLEWARE instead 45 | # Take care to insert it after SessionMiddleware and AuthenticationMiddleware 46 | MIDDLEWARE_CLASSES = ( 47 | 'django.middleware.common.CommonMiddleware', 48 | 'django.contrib.sessions.middleware.SessionMiddleware', 49 | 'django.contrib.auth.middleware.AuthenticationMiddleware', 50 | # ... 51 | 'gae_django_ratelimiter.RateLimiterMiddleware', 52 | # ... 53 | ) 54 | 55 | # ... 56 | 57 | INSTALLED_APPS = ( 58 | 'django.contrib.auth', 59 | 'django.contrib.contenttypes', 60 | 'django.contrib.sessions', 61 | 'django.contrib.sites', 62 | 'django.contrib.messages', 63 | # ... 64 | 'gae_django_ratelimiter', 65 | # ... 66 | ) 67 | ``` 68 | 69 | ### Settings 70 | 71 | - ``GAE_DJANGO_RATELIMITER_ENABLED`` 72 | - Enables rate limiting. Default ``True``. 73 | - ``GAE_DJANGO_RATELIMITER_CACHE_PREFIX`` 74 | - Memcache key prefix. Default ``gaerl``. 75 | - ``GAE_DJANGO_RATELIMITER_CACHE_MINUTES`` 76 | - Interval in minutes. Default ``2``. 77 | - ``GAE_DJANGO_RATELIMITER_CACHE_REQUESTS`` 78 | - Requests per interval. Default ``20``. 79 | - ``GAE_DJANGO_RATELIMITER_COOLDOWN_LAST_REQ`` 80 | - The backoff interval starts from the more recent request after throttling has kicked in. Default ``False``. 81 | - ``GAE_DJANGO_RATELIMITER_COOLDOWN_MINUTES`` 82 | - The backoff interval in minutes if ``GAE_DJANGO_RATELIMITER_COOLDOWN_LAST_REQ`` is True. Equal to ``GAE_DJANGO_RATELIMITER_CACHE_MINUTES`` by default. 83 | - ``GAE_DJANGO_RATELIMITER_INCLUDE_URL_NAMES`` 84 | - List of url names to include. If defined, any url names not in list will be excluded. Default ``[]``. 85 | - ``GAE_DJANGO_RATELIMITER_EXCLUDE_URL_NAMES`` 86 | - List of url names to exclude. If defined, any url names not in list will be included. Default ``[]``. 87 | - ``GAE_DJANGO_RATELIMITER_EXCLUDE_AUTHENTICATED`` 88 | - Exclude authenticated users from limiting. Default ``True``. 89 | - ``GAE_DJANGO_RATELIMITER_EXCLUDE_ADMINS`` 90 | - Exclude admin users from limiting. Default ``True``. 91 | 92 | 93 | ### Advance 94 | 95 | You can subclass the decorator or middleware for your own custom logic. 96 | If using your own middleware, remember to add your own middleware class instead. 97 | 98 | #### Example: Rate limit a crawler bot without affecting other users 99 | 100 | ```python 101 | # myratelimiter.py 102 | import re 103 | import logging 104 | from gae_django_ratelimiter import ratelimit, RateLimiterMiddleware 105 | 106 | 107 | OFFENDING_UA_RE = re.compile(r'^BadBot') 108 | 109 | # Use either a Decorator or Middleware 110 | 111 | # Decorator 112 | class myratelimit(ratelimit): 113 | 114 | def should_ratelimit(self, request): 115 | if not super(myratelimit, self).should_ratelimit(request): 116 | return False 117 | 118 | ua = request.META.get('HTTP_USER_AGENT', '') 119 | return OFFENDING_UA_RE.search(ua) 120 | 121 | def disallowed(self, request): 122 | logging.warn('Too many requests from {}'.format( 123 | request.META.get('REMOTE_ADDR', ''), 124 | ) 125 | return super(myratelimit, self).disallowed(request) 126 | 127 | 128 | # Middleware 129 | class MyRateLimiterMiddleware(RateLimiterMiddleware): 130 | 131 | def should_ratelimit(self, request): 132 | if not super(MyRateLimiterMiddleware, self).should_ratelimit(request): 133 | return False 134 | 135 | ua = request.META.get('HTTP_USER_AGENT', '') 136 | return OFFENDING_UA_RE.search(ua) 137 | ``` 138 | -------------------------------------------------------------------------------- /gae_django_ratelimiter/ratelimiter.py: -------------------------------------------------------------------------------- 1 | 2 | import logging 3 | import functools 4 | from hashlib import md5 5 | import re 6 | from django.http import HttpResponse 7 | from django.core.urlresolvers import resolve 8 | from google.appengine.api import memcache, users 9 | 10 | from .settings import ( 11 | RATELIMITER_ENABLED, RATELIMITER_CACHE_PREFIX, 12 | RATELIMITER_CACHE_MINUTES, RATELIMITER_CACHE_REQUESTS, 13 | RATELIMITER_COOLDOWN_LAST_REQ, RATELIMITER_COOLDOWN_MINUTES, 14 | RATELIMITER_INCLUDE_URL_NAMES, RATELIMITER_EXCLUDE_URL_NAMES, 15 | RATELIMITER_EXCLUDE_AUTHENTICATED, RATELIMITER_EXCLUDE_ADMINS, 16 | ) 17 | 18 | logger = logging.getLogger(__name__) 19 | 20 | 21 | class HttpResponseThrottled(HttpResponse): 22 | status_code = 429 23 | reason_phrase = 'Too Many Requests' 24 | 25 | def __init__(self, content=b'', *args, **kwargs): 26 | super(HttpResponseThrottled, self).__init__( 27 | content, 28 | content_type='text/plain; charset=utf-8', 29 | *args, **kwargs) 30 | if not content: 31 | self.content = '{} {}\nPlease try again later.'.format( 32 | self.status_code, self.reason_phrase 33 | ) 34 | 35 | 36 | class RateLimiter(object): 37 | """Encapsulates the main logic for rate limiting""" 38 | 39 | enabled = RATELIMITER_ENABLED 40 | # The time interval 41 | minutes = RATELIMITER_CACHE_MINUTES 42 | # Number of allowed requests in that interval 43 | requests = RATELIMITER_CACHE_REQUESTS 44 | # Prefix for memcache key 45 | prefix = RATELIMITER_CACHE_PREFIX 46 | # if True, throttling cool down starts after the last req 47 | cooldown_from_last_request = RATELIMITER_COOLDOWN_LAST_REQ 48 | # Cooldown interval in minutes if cooldown_from_last_request is True 49 | cooldown_minutes = RATELIMITER_COOLDOWN_MINUTES 50 | # if False, rate limits also include authenticated users (django/google) 51 | exclude_authenticated = RATELIMITER_EXCLUDE_AUTHENTICATED 52 | # if False, rate limits also include authenticated users (django/google) 53 | exclude_admins = RATELIMITER_EXCLUDE_ADMINS 54 | 55 | # Is mutually exclusive with RateLimiter.exclude_url_names. 56 | # When defined, all urlnames not in include_url_names 57 | # will be excluded. 58 | include_url_names = RATELIMITER_INCLUDE_URL_NAMES 59 | 60 | # Is mutually exclusive with RateLimiter.include_url_names. 61 | # When defined, all urlnames not in include_url_names 62 | # will be included. 63 | exclude_url_names = RATELIMITER_EXCLUDE_URL_NAMES 64 | 65 | # GAE internal IP addresses 66 | # https://cloud.google.com/appengine/docs/standard/python/config/cronref#originating_ip_address 67 | # https://cloud.google.com/appengine/docs/standard/python/taskqueue/push/creating-handlers#writing_a_push_task_request_handler 68 | gae_internal_ips = ('0.1.0.1', '0.1.0.2') 69 | 70 | # Basic bogon ip ranges 71 | bogon_ip_re = re.compile( 72 | r'10\.|127\.|169\.254\.|192\.0\.0\.|192\.168\.|' 73 | '172\.1[6-9]\.|172\.2[0-9]\.|172\.3[0-1]\.') 74 | 75 | def __init__(self, **options): 76 | for key, value in options.items(): 77 | setattr(self, key, value) 78 | 79 | def _is_bogon_ip(self, ip): 80 | if self.bogon_ip_re.match(ip): 81 | return True 82 | return False 83 | 84 | def ip(self, request): 85 | x_forwarded_for = request.META.get('HTTP_X_FORWARDED_FOR') 86 | if x_forwarded_for: 87 | x_forwarded_for_ip = x_forwarded_for.split(',')[-1].strip() 88 | if not (x_forwarded_for_ip in self.gae_internal_ips 89 | or self._is_bogon_ip(x_forwarded_for_ip)): 90 | return x_forwarded_for_ip 91 | return request.META.get('REMOTE_ADDR', '') 92 | 93 | def should_ratelimit(self, request): 94 | if not self.enabled: 95 | return False 96 | 97 | # Skip cron tasks 98 | # https://cloud.google.com/appengine/docs/standard/python/config/cronref#cron_requests 99 | # Safe to use as is because this is protected by GAE 100 | if request.META.get('X-Appengine-Cron', '') == 'true': 101 | return False 102 | 103 | # Skip GAE internal IP addresses 104 | if self.ip(request) in self.gae_internal_ips: 105 | return False 106 | 107 | if self.exclude_admins: 108 | try: 109 | # Django admin 110 | if request.user.is_authenticated and ( 111 | request.user.is_staff or request.user.is_superuser): 112 | return False 113 | # GAE admin 114 | if users.get_current_user() and users.is_current_user_admin(): 115 | return False 116 | except AttributeError as ae: 117 | logger.warning(ae.message) 118 | 119 | try: 120 | if self.exclude_authenticated and ( 121 | request.user.is_authenticated() or 122 | users.get_current_user()): 123 | return False 124 | except AttributeError as ae: 125 | logger.warning(ae.message) 126 | 127 | if self.exclude_url_names: 128 | for url_name in self.exclude_url_names: 129 | if url_name == resolve(request.path_info).url_name: 130 | return False 131 | return True 132 | 133 | if self.include_url_names: 134 | for url_name in self.include_url_names: 135 | if url_name == resolve(request.path_info).url_name: 136 | return True 137 | return False 138 | return True 139 | 140 | def disallowed(self, request): 141 | """Override this method if you want to log incidents""" 142 | return HttpResponseThrottled() 143 | 144 | def current_key(self, request): 145 | """Override this to use a different cache key""" 146 | # Google's memcache key len is max 250 bytes 147 | # https://cloud.google.com/appengine/docs/standard/python/memcache/ 148 | m = md5() 149 | # Use a basic hash of the UA 150 | m.update(request.META.get('HTTP_USER_AGENT', '')) 151 | return '{}_{}_{}_{}'.format( 152 | self.prefix, 153 | self.ip(request), 154 | m.hexdigest(), 155 | self.minutes, 156 | ) 157 | 158 | def expire_after(self): 159 | """Used for setting the memcache expiry""" 160 | return (self.minutes) * 60 161 | 162 | def cached_count(self, key): 163 | return memcache.get(key, 0) or 0 164 | 165 | def cache_incr(self, key): 166 | # add first, to ensure the key exists 167 | added = memcache.add(key, 0, time=self.expire_after()) 168 | if not added and self.cooldown_from_last_request: 169 | # already exists so we extend the memcache expiry 170 | # by another interval 171 | memcache.set( 172 | key, self.cached_count(key), time=self.cooldown_minutes * 60) 173 | memcache.incr(key) 174 | 175 | def check(self, request): 176 | if not self.should_ratelimit(request): 177 | return None, 0 178 | 179 | # Increment rate limiting counter 180 | self.cache_incr(self.current_key(request)) 181 | cached_count = self.cached_count(self.current_key(request)) 182 | 183 | if cached_count > self.requests: 184 | return self.disallowed(request), cached_count 185 | 186 | return None, cached_count 187 | 188 | 189 | # Middleware 190 | class RateLimiterMiddleware(RateLimiter): 191 | 192 | def __init__(self, get_response=None, **options): 193 | self.get_response = get_response 194 | super(RateLimiterMiddleware, self).__init__(**options) 195 | 196 | # For Django>=1.10 197 | # https://docs.djangoproject.com/en/1.11/topics/http/middleware/#upgrading-pre-django-1-10-style-middleware 198 | def __call__(self, request): 199 | response = self.process_request(request) 200 | if not response: 201 | response = self.get_response(request) 202 | response = self.process_response(request, response) 203 | return response 204 | 205 | def process_request(self, request): 206 | res, _ = self.check(request) 207 | return res 208 | 209 | def process_response(self, request, response): 210 | if response.status_code < 400 and self.should_ratelimit(request): 211 | response['X-Rate-Limit-Remaining-{}'.format(self.minutes)] = ( 212 | self.requests - self.cached_count(self.current_key(request))) 213 | return response 214 | 215 | 216 | # Decorator 217 | # Modified from 218 | # https://github.com/simonw/ratelimitcache/blob/master/ratelimitcache.py 219 | class ratelimit(RateLimiter): 220 | 221 | def __call__(self, fn): 222 | def wrapper(request, *args, **kwargs): 223 | return self.view_wrapper(request, fn, *args, **kwargs) 224 | functools.update_wrapper(wrapper, fn) 225 | return wrapper 226 | 227 | def view_wrapper(self, request, fn, *args, **kwargs): 228 | res, cached_count = self.check(request) 229 | if res: 230 | return res 231 | 232 | res = fn(request, *args, **kwargs) 233 | if cached_count: 234 | res['X-Rate-Limit-Remaining-{}'.format(self.minutes)] = ( 235 | self.requests - cached_count) 236 | return res 237 | -------------------------------------------------------------------------------- /tests/test_ratelimiter.py: -------------------------------------------------------------------------------- 1 | 2 | import sys 3 | import unittest 4 | import copy 5 | import time 6 | from hashlib import md5 7 | try: 8 | import unittest.mock as compat_mock 9 | except ImportError: 10 | import mock as compat_mock 11 | 12 | from django.test import Client 13 | from django.core.exceptions import ImproperlyConfigured 14 | try: # pragma: no cover 15 | from google.appengine.ext import testbed 16 | except ImportError: 17 | sys.path.insert(1, 'gae_sdk/google_appengine') 18 | sys.path.insert(1, 'gae_sdk/google_appengine/lib/yaml/lib/') 19 | from google.appengine.ext import testbed 20 | 21 | from django.conf import settings 22 | from django import setup as django_setup 23 | 24 | required_middleware = [ 25 | 'django.contrib.sessions.middleware.SessionMiddleware', 26 | 'django.contrib.auth.middleware.AuthenticationMiddleware', 27 | ] 28 | 29 | try: 30 | from gae_django_ratelimiter import RateLimiterMiddleware 31 | from gae_django_ratelimiter.ratelimiter import RateLimiter 32 | from middleware import TestRateLimiterMiddleware 33 | except ImproperlyConfigured: 34 | settings.configure() 35 | settings.ALLOWED_HOSTS = ['testserver'] 36 | settings.ROOT_URLCONF = 'urls' 37 | settings.MIDDLEWARE_CLASSES = required_middleware 38 | settings.INSTALLED_APPS = [ 39 | 'django.contrib.sessions', 40 | 'django.contrib.auth', 41 | 'django.contrib.contenttypes', 42 | 'tests', 43 | ] 44 | from gae_django_ratelimiter import RateLimiterMiddleware 45 | from gae_django_ratelimiter.ratelimiter import RateLimiter 46 | from middleware import TestRateLimiterMiddleware 47 | 48 | django_setup() 49 | 50 | 51 | class RateLimiterTests(unittest.TestCase): 52 | 53 | def setUp(self): 54 | self.testbed = testbed.Testbed() 55 | self.testbed.activate() 56 | self.testbed.init_memcache_stub() 57 | self.testbed.init_app_identity_stub() 58 | 59 | settings.MIDDLEWARE_CLASSES = ( 60 | required_middleware + ['middleware.TestRateLimiterMiddleware']) 61 | 62 | self.request = compat_mock.Mock() 63 | self.request.META = { 64 | 'REMOTE_ADDR': '127.0.0.1', 65 | 'HTTP_USER_AGENT': 'Mozilla/5.0 RATELIMITER', 66 | } 67 | self.request.url_name = 'random' 68 | 69 | def tearDown(self): 70 | self.testbed.deactivate() 71 | 72 | def test_basic(self): 73 | req = copy.deepcopy(self.request) 74 | 75 | test_requests = 9 76 | test_minutes = 1 77 | test_prefix = 'xyz' 78 | rl = RateLimiter( 79 | requests=test_requests, minutes=test_minutes, prefix=test_prefix) 80 | 81 | self.assertEqual(test_requests, rl.requests) 82 | self.assertEqual(test_minutes, rl.minutes) 83 | self.assertEqual(test_prefix, rl.prefix) 84 | 85 | self.assertEquals(req.META['REMOTE_ADDR'], rl.ip(req)) 86 | req.META['HTTP_X_FORWARDED_FOR'] = '::1' 87 | self.assertEquals(req.META['HTTP_X_FORWARDED_FOR'], rl.ip(req)) 88 | 89 | m = md5() 90 | m.update(req.META.get('HTTP_USER_AGENT', '')) 91 | self.assertEqual( 92 | '{}_{}_{}_{}'.format('xyz', rl.ip(req), m.hexdigest(), rl.minutes), 93 | rl.current_key(req)) 94 | 95 | def test_disabled(self): 96 | from middleware import DisabledRateLimiterMiddleware 97 | settings.MIDDLEWARE_CLASSES = ( 98 | required_middleware + ['middleware.DisabledRateLimiterMiddleware']) 99 | 100 | c = Client() 101 | res = c.get('/random') 102 | self.assertEqual(200, res.status_code) 103 | self.assertIsNone( 104 | res.get('X-Rate-Limit-Remaining-{}'.format( 105 | DisabledRateLimiterMiddleware.minutes))) 106 | 107 | def test_gae_internal(self): 108 | c = Client() 109 | 110 | gae_ip = '0.1.0.1' 111 | res = c.get('/random', **{'REMOTE_ADDR': gae_ip}) 112 | self.assertEqual(200, res.status_code) 113 | self.assertIsNone( 114 | res.get('X-Rate-Limit-Remaining-{}'.format( 115 | TestRateLimiterMiddleware.requests))) 116 | 117 | res = c.get('/random', **{'X-Appengine-Cron': 'true'}) 118 | self.assertEqual(200, res.status_code) 119 | self.assertIsNone( 120 | res.get('X-Rate-Limit-Remaining-{}'.format( 121 | TestRateLimiterMiddleware.requests))) 122 | 123 | res = c.get( 124 | '/random', 125 | **{'REMOTE_ADDR': '192.168.0.1', 'HTTP_X_FORWARDED_FOR': gae_ip}) 126 | self.assertEqual(200, res.status_code) 127 | self.assertIsNotNone( 128 | res.get('X-Rate-Limit-Remaining-{}'.format( 129 | TestRateLimiterMiddleware.minutes))) 130 | 131 | def test_xforwardedfor(self): 132 | req = copy.deepcopy(self.request) 133 | req.META['HTTP_X_FORWARDED_FOR'] = '::1' 134 | rl = RateLimiter() 135 | self.assertEquals(req.META['HTTP_X_FORWARDED_FOR'], rl.ip(req)) 136 | 137 | # bogon IP 138 | req.META['HTTP_X_FORWARDED_FOR'] = '10.0.0.1' 139 | self.assertEquals(req.META['REMOTE_ADDR'], rl.ip(req)) 140 | 141 | def test_cooldownfromlastreq(self): 142 | 143 | from middleware import CooldownRateLimiterMiddleware 144 | settings.MIDDLEWARE_CLASSES = ( 145 | required_middleware + ['middleware.CooldownRateLimiterMiddleware']) 146 | 147 | c = Client() 148 | window_time = CooldownRateLimiterMiddleware.minutes * 60 149 | 150 | for i in range(CooldownRateLimiterMiddleware.requests): 151 | res = c.get('/random') 152 | self.assertEqual(200, res.status_code) 153 | time.sleep(window_time // CooldownRateLimiterMiddleware.requests) 154 | 155 | res = c.get('/random') 156 | self.assertEqual(429, res.status_code) 157 | # wait till cache expires 158 | time.sleep(CooldownRateLimiterMiddleware.cooldown_minutes * 60) 159 | res = c.get('/random') 160 | self.assertEqual(200, res.status_code) 161 | 162 | def test_include_url_names(self): 163 | settings.MIDDLEWARE_CLASSES = ( 164 | required_middleware + ['middleware.IncludeRateLimiterMiddleware']) 165 | c = Client() 166 | res_include = c.get('/random') 167 | res_exclude = c.get('/notrandom') 168 | self.assertIsNotNone( 169 | res_include.get('X-Rate-Limit-Remaining-{}'.format( 170 | RateLimiterMiddleware.minutes))) 171 | self.assertIsNone( 172 | res_exclude.get('X-Rate-Limit-Remaining-{}'.format( 173 | RateLimiterMiddleware.minutes))) 174 | 175 | def test_exclude_url_names(self): 176 | settings.MIDDLEWARE_CLASSES = ( 177 | required_middleware + ['middleware.ExcludeRateLimiterMiddleware']) 178 | 179 | c = Client() 180 | res_include = c.get('/notrandom') 181 | res_exclude = c.get('/random') 182 | self.assertIsNotNone( 183 | res_include.get('X-Rate-Limit-Remaining-{}'.format( 184 | RateLimiterMiddleware.minutes))) 185 | self.assertIsNone( 186 | res_exclude.get('X-Rate-Limit-Remaining-{}'.format( 187 | RateLimiterMiddleware.minutes))) 188 | 189 | def _loginUser(self, email='user@example.com', id='123', is_admin=False): 190 | self.testbed.setup_env( 191 | user_email=email, 192 | user_id=id, 193 | user_is_admin='1' if is_admin else '0', 194 | overwrite=True) 195 | 196 | def test_admin_user(self): 197 | from middleware import ( 198 | ExcludeAuthAdminRateLimiterMiddleware, 199 | ) 200 | middlewares = [ 201 | ExcludeAuthAdminRateLimiterMiddleware, 202 | ] 203 | 204 | for m in middlewares: 205 | settings.MIDDLEWARE_CLASSES = ( 206 | required_middleware + ['{}.{}'.format( 207 | m.__module__, m.__name__ 208 | )]) 209 | 210 | c = Client() 211 | res = c.get('/random') 212 | self.assertEqual(200, res.status_code) 213 | self.assertIsNotNone( 214 | res.get('X-Rate-Limit-Remaining-{}'.format( 215 | m.minutes)), '{} failed'.format(m.__name__)) 216 | 217 | self._loginUser(is_admin=True) 218 | res = c.get('/random') 219 | self.assertEqual(200, res.status_code) 220 | if m.exclude_admins: 221 | self.assertIsNone( 222 | res.get('X-Rate-Limit-Remaining-{}'.format( 223 | m.minutes)), '{} failed'.format(m.__name__)) 224 | else: 225 | self.assertIsNotNone( 226 | res.get('X-Rate-Limit-Remaining-{}'.format( 227 | m.minutes)), '{} failed'.format(m.__name__)) 228 | 229 | def test_authenticated_user(self): 230 | from middleware import ( 231 | ExcludeAuthRateLimiterMiddleware, 232 | IncludeAuthRateLimiterMiddleware, 233 | ExcludeAuthAdminRateLimiterMiddleware, 234 | ) 235 | middlewares = [ 236 | ExcludeAuthRateLimiterMiddleware, 237 | IncludeAuthRateLimiterMiddleware, 238 | ExcludeAuthAdminRateLimiterMiddleware, 239 | ] 240 | 241 | for m in middlewares: 242 | settings.MIDDLEWARE_CLASSES = ( 243 | required_middleware + ['{}.{}'.format( 244 | m.__module__, m.__name__ 245 | )]) 246 | 247 | c = Client() 248 | res = c.get('/random') 249 | self.assertEqual(200, res.status_code) 250 | self.assertIsNotNone( 251 | res.get('X-Rate-Limit-Remaining-{}'.format( 252 | m.minutes)), '{} failed'.format(m.__name__)) 253 | 254 | self._loginUser() 255 | res = c.get('/random') 256 | self.assertEqual(200, res.status_code) 257 | if m.exclude_authenticated: 258 | self.assertIsNone( 259 | res.get('X-Rate-Limit-Remaining-{}'.format( 260 | m.minutes)), '{} failed'.format(m.__name__)) 261 | else: 262 | self.assertIsNotNone( 263 | res.get('X-Rate-Limit-Remaining-{}'.format( 264 | m.minutes)), '{} failed'.format(m.__name__)) 265 | 266 | 267 | class RateLimiterMiddlewareTests(unittest.TestCase): 268 | 269 | def setUp(self): 270 | self.testbed = testbed.Testbed() 271 | self.testbed.activate() 272 | self.testbed.init_memcache_stub() 273 | 274 | def tearDown(self): 275 | self.testbed.deactivate() 276 | 277 | def test_basic(self): 278 | settings.MIDDLEWARE_CLASSES = ( 279 | required_middleware + ['middleware.TestRateLimiterMiddleware']) 280 | 281 | c = Client() 282 | for i in range(1, TestRateLimiterMiddleware.requests + 2): 283 | res = c.get('/random') 284 | if i <= TestRateLimiterMiddleware.requests: 285 | self.assertEqual( 286 | TestRateLimiterMiddleware.requests - i, 287 | int(res.get('X-Rate-Limit-Remaining-{}'.format( 288 | TestRateLimiterMiddleware.minutes), '-1'))) 289 | else: 290 | self.assertIsNone( 291 | res.get('X-Rate-Limit-Remaining-{}'.format( 292 | TestRateLimiterMiddleware.requests))) 293 | 294 | def test_django1_10(self): 295 | settings.MIDDLEWARE_CLASSES = ( 296 | required_middleware + ['middleware.TestRateLimiterMiddleware']) 297 | 298 | c = Client() 299 | res = c.get('/random') 300 | self.assertEqual( 301 | TestRateLimiterMiddleware.requests - 1, 302 | int(res.get('X-Rate-Limit-Remaining-{}'.format( 303 | TestRateLimiterMiddleware.minutes)))) 304 | 305 | 306 | class RateLimitDecoratorTests(unittest.TestCase): 307 | 308 | def setUp(self): 309 | self.testbed = testbed.Testbed() 310 | self.testbed.activate() 311 | self.testbed.init_memcache_stub() 312 | 313 | def tearDown(self): 314 | self.testbed.deactivate() 315 | 316 | def test_basic(self): 317 | # settings.MIDDLEWARE_CLASSES = required_middleware 318 | 319 | requests_limit = 5 320 | minutes = 1 321 | c = Client() 322 | for i in range(1, requests_limit + 2): 323 | res = c.get('/decorated') 324 | 325 | if i <= requests_limit: 326 | self.assertEqual( 327 | requests_limit - i, 328 | int(res.get('X-Rate-Limit-Remaining-{}'.format( 329 | minutes), '-1')) 330 | ) 331 | else: 332 | self.assertIsNone( 333 | res.get('X-Rate-Limit-Remaining-{}'.format(minutes))) 334 | --------------------------------------------------------------------------------