Web auth

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

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,82 +19,118 @@
 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")
         except Exception as e:
             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']

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,18 +86,21 @@ 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'}))
     response.headers["Content-type"] = "application/json"
     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)