wsgi: Add middleware for delaying and rejecting requests

- Fix #7

kw/platform/settings.py

@@ -0,0 +1,6 @@
+VALID_USER_AGENT_RE = (
+    r"^(?P<name>\S.+?)\/(?P<version>\S.+?) "
+    r"\(Kiwi\.com (?P<environment>\S.+?)\)(?: ?(?P<system_info>.*))$"
+)
+REQUESTS_SLOWDOWN_DATETIME = "2019-06-24T13:00:00"
+REQUESTS_RESTRICT_DATETIME = "2019-08-01T13:00:00"

kw/platform/wsgi.py

@@ -0,0 +1,78 @@
+import re
+import time
+from datetime import datetime
+
+import webob
+from dateutil.parser import parse
+from webob.dec import wsgify
+
+from . import settings
+
+
+user_agent_re = re.compile(settings.VALID_USER_AGENT_RE)
+req_slowdown_datetime = parse(settings.REQUESTS_SLOWDOWN_DATETIME)
+req_restrict_datetime = parse(settings.REQUESTS_RESTRICT_DATETIME)
+
+
+def _refuse_request(req, app):
+    return webob.exc.HTTPBadRequest("Bad User-Agent: see Kiwi.com RFC #22")
+
+
+def _slowdown_request(req, app):
+    before_time = time.time()
+    resp = req.get_response(app)
+    seconds = time.time() - before_time
+    time.sleep(seconds)
+    return resp
+
+
+@wsgify.middleware
+def user_agent_middleware(req, app):
+    """Validate client's User-Agent header and modify response based on that.
+
+    If the User-Agent header is invalid, there are three things that can happen:
+
+    1. The current time is less then :obj:`settings.REQUESTS_SLOWDOWN_DATETIME`,
+       do nothing in this case.
+    2. The current time is less then :obj:`settings.REQUESTS_RESTRICT_DATETIME`,
+       slow down the response twice the normal responce time.
+    3. The current time is more then :obj:`settings.REQUESTS_RESTRICT_DATETIME`,
+       refuse the request, return HTTP 400 to the client.
+
+    Usage:
+
+        from your_app import wsgi_app
+
+        wsgi_app = user_agent_middleware(wsgi_app)
+
+    For example, in Flask, the middleware can be applied like this:
+
+        from flask import Flask
+
+        app = Flask(__name__)
+        app.wsgi_app = user_agent_middleware(app)
+
+        app.run()
+
+    In Django, you can apply the middleware like this:
+
+        from django.core.wsgi import get_wsgi_application
+
+        application = user_agent_middleware(get_wsgi_application())
+
+    For more information see
+    `Django docs <https://docs.djangoproject.com/en/dev/howto/deployment/wsgi/>`_.
+    """
+    now = datetime.utcnow()
+
+    if now < req_slowdown_datetime:
+        return app
+
+    is_valid = bool(user_agent_re.match(req.user_agent))
+
+    if is_valid:
+        return app
+    elif now < req_restrict_datetime:
+        return _slowdown_request(req, app)
+    else:
+        return _refuse_request(req, app)

pyproject.toml

@@ -42,7 +42,10 @@ include = ["*.md", "*.toml", "*.txt", "*.yaml", ".coveragerc", "tox.ini"]
 
 [tool.poetry.dependencies]
 python = "~2.7 || ^3.5"
+webob = "^1.8"
+python-dateutil = "^2.8"
 
 [tool.poetry.dev-dependencies]
 pytest = "^4.6"
+pytest-mock = "^1.10"
 coverage = "^4.5"

test/test_wsgi.py

@@ -0,0 +1,47 @@
+import time
+from datetime import datetime
+
+import pytest
+from webob.request import BaseRequest
+
+from kw.platform import wsgi as uut
+
+
+def create_app(sleep_seconds=0):
+    def simple_app(environ, start_response):
+        if sleep_seconds:
+            time.sleep(sleep_seconds)
+        start_response("200 OK", [("Content-Type", "text/html; charset=UTF-8")])
+        return ["OK"]
+
+    return simple_app
+
+
+@pytest.mark.parametrize(
+    "user_agent,sleep_seconds,request_time,expected_status,current_time",
+    [
+        ("invalid", 0, 0, 200, datetime(2019, 5, 21, 0, 0, 0)),
+        ("invalid", 0, 0, 400, datetime(2020, 1, 10, 0, 0, 0)),
+        ("mambo/1a (Kiwi.com dev)", 0, 0, 200, datetime(2020, 1, 10, 0, 0, 0)),
+        ("invalid", 0.1, 0.2, 200, datetime(2019, 7, 1, 0, 0, 0)),
+        ("mambo/1a (Kiwi.com dev)", 0.1, 0.1, 200, datetime(2019, 7, 1, 0, 0, 0)),
+    ],
+)
+def test_user_agent_middleware(
+    mocker, user_agent, sleep_seconds, request_time, expected_status, current_time
+):
+    req = BaseRequest.blank("/")
+    req.user_agent = user_agent
+
+    mocker.patch(
+        "kw.platform.wsgi.datetime",
+        mocker.Mock(utcnow=mocker.Mock(return_value=current_time)),
+    )
+    app = uut.user_agent_middleware(create_app(sleep_seconds))
+
+    before_time = time.time()
+    res = req.get_response(app)
+    request_time = time.time() - before_time
+
+    assert res.status_code == expected_status
+    assert request_time >= request_time

tox.ini

@@ -40,7 +40,6 @@ atomic = true
 force_grid_wrap = 0
 include_trailing_comma = true
 lines_after_imports = 2
-lines_between_types = 1
 multi_line_output = 3
 not_skip = __init__.py
 use_parentheses = true