From cab8aed6be543f045d268fd41110976f43fb13e1 Mon Sep 17 00:00:00 2001 From: Sandy Chapman Date: Thu, 21 May 2015 09:55:35 -0300 Subject: [PATCH 1/2] Added the ability to enable Basic web authentication via specifying a LOCUST_USERNAME and LOCUST_PASSWORD environment variable pair. --- locust/web.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/locust/web.py b/locust/web.py index 7887a01364..f4d99f1208 100644 --- a/locust/web.py +++ b/locust/web.py @@ -16,6 +16,8 @@ from .runners import MasterLocustRunner from locust.stats import median_from_dict from locust import version +from functools import wraps +from flask import request, Response import logging logger = logging.getLogger(__name__) @@ -27,7 +29,35 @@ app.root_path = os.path.dirname(os.path.abspath(__file__)) +def check_auth(username, password): + """This function is called to check if a username / + password combination is valid. + """ + try: + locust_username = os.environ['LOCUST_USER_NAME'] + locust_password = os.environ['LOCUST_PASSWORD'] + return username == locust_username and password == locust_password + except: + return True + +def authenticate(): + """Sends a 401 response that enables basic auth""" + return Response( + 'Could not verify your access level for that URL.\n' + 'You have to login with proper credentials', 401, + {'WWW-Authenticate': 'Basic realm="Login Required"'}) + +def requires_auth(f): + @wraps(f) + def decorated(*args, **kwargs): + auth = request.authorization + if (auth == None and not check_auth('', '')) or (auth != None and not check_auth(auth.username, auth.password)): + return authenticate() + return f(*args, **kwargs) + return decorated + @app.route('/') +@requires_auth def index(): is_distributed = isinstance(runners.locust_runner, MasterLocustRunner) if is_distributed: @@ -44,6 +74,7 @@ def index(): ) @app.route('/swarm', methods=["POST"]) +@requires_auth def swarm(): assert request.method == "POST" @@ -55,6 +86,7 @@ def swarm(): return response @app.route('/stop') +@requires_auth def stop(): runners.locust_runner.stop() response = make_response(json.dumps({'success':True, 'message': 'Test stopped'})) @@ -62,11 +94,13 @@ def stop(): return response @app.route("/stats/reset") +@requires_auth def reset_stats(): runners.locust_runner.stats.reset_all() return "ok" @app.route("/stats/requests/csv") +@requires_auth def request_stats_csv(): rows = [ ",".join([ @@ -105,6 +139,7 @@ def request_stats_csv(): return response @app.route("/stats/distribution/csv") +@requires_auth def distribution_stats_csv(): rows = [",".join(( '"Name"', @@ -133,6 +168,7 @@ def distribution_stats_csv(): return response @app.route('/stats/requests') +@requires_auth @memoize(timeout=DEFAULT_CACHE_TIME, dynamic_timeout=True) def request_stats(): stats = [] @@ -175,12 +211,14 @@ def request_stats(): return json.dumps(report) @app.route("/exceptions") +@requires_auth def exceptions(): response = make_response(json.dumps({'exceptions': [{"count": row["count"], "msg": row["msg"], "traceback": row["traceback"], "nodes" : ", ".join(row["nodes"])} for row in runners.locust_runner.exceptions.itervalues()]})) response.headers["Content-type"] = "application/json" return response @app.route("/exceptions/csv") +@requires_auth def exceptions_csv(): data = StringIO() writer = csv.writer(data) From 0f63a830c6ec4f684883f90040f21491dfa83346 Mon Sep 17 00:00:00 2001 From: Sandy Chapman Date: Wed, 27 May 2015 14:59:23 -0300 Subject: [PATCH 2/2] Added documentation and test coverage for web auth. --- docs/extending-locust.rst | 46 ++++++++++++++++++++---- locust/test/test_web.py | 74 +++++++++++++++++++++++++++++---------- 2 files changed, 96 insertions(+), 24 deletions(-) diff --git a/docs/extending-locust.rst b/docs/extending-locust.rst index ca2d7aa3d4..ab7f7f7fa3 100644 --- a/docs/extending-locust.rst +++ b/docs/extending-locust.rst @@ -7,16 +7,16 @@ Locust comes with a number of events that provides hooks for extending locust in Event listeners can be registered at the module level in a locust file. Here's an example:: from locust import events - + def my_success_handler(method, path, response_time, response, **kw): print "Successfully fetched: %s" % (path) - + events.request_success += my_success_handler .. note:: - It's highly recommended that you add a wildcard keyword argument in your listeners - (the \**kw in the code above), to prevent your code from breaking if new arguments are + It's highly recommended that you add a wildcard keyword argument in your listeners + (the \**kw in the code above), to prevent your code from breaking if new arguments are added in a future version. .. seealso:: @@ -28,13 +28,47 @@ Event listeners can be registered at the module level in a locust file. Here's a Adding Web Routes ================== -Locust uses Flask to serve the web UI and therefore it is easy to add web end-points to the web UI. +Locust uses Flask to serve the web UI and therefore it is easy to add web end-points to the web UI. Just import the Flask app in your locustfile and set up a new route:: from locust import web - + @web.app.route("/added_page") def my_added_page(): return "Another page" You should now be able to start locust and browse to http://127.0.0.1:8089/added_page + + +Enable Authenticated Access +=========================== + +.. attention:: + + Using this feature in no way makes the data transfered between locusts or between + the locust master and the client secure. + +Locust supports basic access authentication. This is a simply form of authentication +where a browser will prompt for a username and password if a HTTP 401 response is +returned with a special header. + +The idea behind using this sort of authentication is not for securing the data transferred +between the client and locust, as that is still transmitted in plain text, however, it will +prevent unauthorized tampering via the web interface with a test. This could also be useful +if a load tester is used in a multi-developer environment where each developer could invoke +locust with their own credentials. + +Enabling support for basic authentication is done by specifying two environment variables. + +You can put the following lines in your Bash `.profile` or `.bashrc`. + +.. code:: + + export LOCUST_USER_NAME=my_user + export LOCUST_PASSWORD=secret + +Optionally, you can also set them prior to invoking locust. + +.. code:: + + LOCUST_USER_NAME=my_user LOCUST_PASSWORD=secret locust diff --git a/locust/test/test_web.py b/locust/test/test_web.py index dad15fc412..71733a011c 100644 --- a/locust/test/test_web.py +++ b/locust/test/test_web.py @@ -2,6 +2,8 @@ import json import sys import traceback +import os +from base64 import b64encode from StringIO import StringIO import requests @@ -17,65 +19,65 @@ class TestWebUI(LocustTestCase): def setUp(self): super(TestWebUI, self).setUp() - + stats.global_stats.clear_all() parser = parse_options()[0] options = parser.parse_args([])[0] runners.locust_runner = LocustRunner([], options) - + self._web_ui_server = wsgi.WSGIServer(('127.0.0.1', 0), web.app, log=None) gevent.spawn(lambda: self._web_ui_server.serve_forever()) gevent.sleep(0.01) self.web_port = self._web_ui_server.server_port - + def tearDown(self): super(TestWebUI, self).tearDown() self._web_ui_server.stop() - + def test_index(self): self.assertEqual(200, requests.get("http://127.0.0.1:%i/" % self.web_port).status_code) - + def test_stats_no_data(self): self.assertEqual(200, requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port).status_code) - + def test_stats(self): stats.global_stats.get("/test", "GET").log(120, 5612) response = requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port) self.assertEqual(200, response.status_code) - + data = json.loads(response.content) self.assertEqual(2, len(data["stats"])) # one entry plus Total self.assertEqual("/test", data["stats"][0]["name"]) self.assertEqual("GET", data["stats"][0]["method"]) self.assertEqual(120, data["stats"][0]["avg_response_time"]) - + def test_stats_cache(self): stats.global_stats.get("/test", "GET").log(120, 5612) response = requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port) self.assertEqual(200, response.status_code) data = json.loads(response.content) self.assertEqual(2, len(data["stats"])) # one entry plus Total - + # add another entry stats.global_stats.get("/test2", "GET").log(120, 5612) data = json.loads(requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port).content) self.assertEqual(2, len(data["stats"])) # old value should be cached now - + web.request_stats.clear_cache() - + data = json.loads(requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port).content) self.assertEqual(3, len(data["stats"])) # this should no longer be cached - + def test_request_stats_csv(self): stats.global_stats.get("/test", "GET").log(120, 5612) response = requests.get("http://127.0.0.1:%i/stats/requests/csv" % self.web_port) self.assertEqual(200, response.status_code) - + def test_distribution_stats_csv(self): stats.global_stats.get("/test", "GET").log(120, 5612) response = requests.get("http://127.0.0.1:%i/stats/distribution/csv" % self.web_port) self.assertEqual(200, response.status_code) - + def test_exceptions_csv(self): try: raise Exception("Test exception") @@ -83,16 +85,52 @@ def test_exceptions_csv(self): tb = sys.exc_info()[2] runners.locust_runner.log_exception("local", str(e), "".join(traceback.format_tb(tb))) runners.locust_runner.log_exception("local", str(e), "".join(traceback.format_tb(tb))) - + response = requests.get("http://127.0.0.1:%i/exceptions/csv" % self.web_port) self.assertEqual(200, response.status_code) - + reader = csv.reader(StringIO(response.content)) rows = [] for row in reader: rows.append(row) - + self.assertEqual(2, len(rows)) self.assertEqual("Test exception", rows[1][1]) self.assertEqual(2, int(rows[1][0]), "Exception count should be 2") - + + def test_web_auth_unauthorized(self): + os.environ['LOCUST_USER_NAME'] = 'my_user' + os.environ['LOCUST_PASSWORD'] = 'supersecret' + + response = requests.get("http://127.0.0.1:%i/" % self.web_port) + self.assertEqual(401, response.status_code) + + del os.environ['LOCUST_USER_NAME'] + del os.environ['LOCUST_PASSWORD'] + + def test_web_auth_authorized(self): + os.environ['LOCUST_USER_NAME'] = 'my_user' + os.environ['LOCUST_PASSWORD'] = 'supersecret' + + auth_value = 'Basic ' + b64encode('my_user:supersecret') + + response = requests.get("http://127.0.0.1:%i/" % self.web_port, headers={'Authorization': auth_value}) + self.assertEqual(200, response.status_code) + + del os.environ['LOCUST_USER_NAME'] + del os.environ['LOCUST_PASSWORD'] + + def test_web_auth_both(self): + os.environ['LOCUST_USER_NAME'] = 'my_user' + os.environ['LOCUST_PASSWORD'] = 'supersecret' + + response = requests.get("http://127.0.0.1:%i/" % self.web_port) + self.assertEqual(401, response.status_code) + + auth_value = 'Basic ' + b64encode('my_user:supersecret') + + response = requests.get("http://127.0.0.1:%i/" % self.web_port, headers={'Authorization': auth_value}) + self.assertEqual(200, response.status_code) + + del os.environ['LOCUST_USER_NAME'] + del os.environ['LOCUST_PASSWORD']