Merge pull request #1354 from mattdodge/tls-certs

Add CLI params for TLS cert and key - serves over HTTPS

locust/argument_parser.py

@@ -178,6 +178,16 @@ def setup_parser_arguments(parser):
         default=None,
         help='Turn on Basic Auth for the web interface. Should be supplied in the following format: username:password'
     )
+    web_ui_group.add_argument(
+        '--tls-cert',
+        default="",
+        help="Optional path to TLS certificate to use to serve over HTTPS"
+    )
+    web_ui_group.add_argument(
+        '--tls-key',
+        default="",
+        help="Optional path to TLS private key to use to serve over HTTPS"
+    )
 
     master_group = parser.add_argument_group(
         "Master options", 

locust/env.py

@@ -116,14 +116,18 @@ def create_worker_runner(self, master_host, master_port):
             master_port=master_port,
         )
 
-    def create_web_ui(self, host="", port=8089, auth_credentials=None):
+    def create_web_ui(self, host="", port=8089, auth_credentials=None, tls_cert=None, tls_key=None):
         """
         Creates a :class:`WebUI <locust.web.WebUI>` instance for this Environment and start running the web server
         
         :param host: Host/interface that the web server should accept connections to. Defaults to ""
                      which means all interfaces
         :param port: Port that the web server should listen to
         :param auth_credentials: If provided (in format "username:password") basic auth will be enabled
+        :param tls_cert: An optional path (str) to a TLS cert. If this is provided the web UI will be
+                         served over HTTPS
+        :param tls_key: An optional path (str) to a TLS private key. If this is provided the web UI will be
+                        served over HTTPS
         """
-        self.web_ui = WebUI(self, host, port, auth_credentials=auth_credentials)
+        self.web_ui = WebUI(self, host, port, auth_credentials=auth_credentials, tls_cert=tls_cert, tls_key=tls_key)
         return self.web_ui

locust/main.py

@@ -224,14 +224,16 @@ def timelimit_stop():
     # start Web UI
     if not options.headless and not options.worker:
         # spawn web greenlet
-        logger.info("Starting web monitor at http://%s:%s" % (options.web_host, options.web_port))
+        protocol = "https" if options.tls_cert and options.tls_key else "http"
+        logger.info("Starting web monitor at %s://%s:%s" % (protocol, options.web_host, options.web_port))
         try:
             if options.web_host == "*":
                 # special check for "*" so that we're consistent with --master-bind-host
                 web_host = ''
             else:
                 web_host = options.web_host
-            web_ui = environment.create_web_ui(host=web_host, port=options.web_port, auth_credentials=options.web_auth)
+            web_ui = environment.create_web_ui(
+                host=web_host, port=options.web_port, auth_credentials=options.web_auth, tls_cert=options.tls_cert, tls_key=options.tls_key)
         except AuthCredentialsError:
             logger.error("Credentials supplied with --web-auth should have the format: username:password")
             sys.exit(1)

locust/test/test_web.py

@@ -1,9 +1,12 @@
 # -*- coding: utf-8 -*-
 import csv
 import json
+import os
 import sys
 import traceback
+from datetime import datetime, timedelta
 from io import StringIO
+from tempfile import NamedTemporaryFile
 
 import gevent
 import requests
@@ -305,3 +308,71 @@ def test_index_with_basic_auth_enabled_incorrect_credentials(self):
 
     def test_index_with_basic_auth_enabled_blank_credentials(self):
         self.assertEqual(401, requests.get("http://127.0.0.1:%i/?ele=phino" % self.web_port).status_code)
+
+
+class TestWebUIWithTLS(LocustTestCase):
+
+    def _create_tls_cert(self):
+        """ Generate a TLS cert and private key to serve over https """
+        from cryptography import x509
+        from cryptography.x509.oid import NameOID
+        from cryptography.hazmat.primitives import hashes
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives import serialization
+        from cryptography.hazmat.primitives.asymmetric import rsa
+
+        key = rsa.generate_private_key(public_exponent=2**16+1, key_size=2048, backend=default_backend())
+        name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "127.0.0.1")])
+        now = datetime.utcnow()
+        cert = (
+            x509.CertificateBuilder()
+            .subject_name(name)
+            .issuer_name(name)
+            .public_key(key.public_key())
+            .serial_number(1000)
+            .not_valid_before(now)
+            .not_valid_after(now + timedelta(days=10*365))
+            .sign(key, hashes.SHA256(), default_backend())
+        )
+        cert_pem = cert.public_bytes(encoding=serialization.Encoding.PEM)
+        key_pem = key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+        return cert_pem, key_pem
+
+    def setUp(self):
+        super(TestWebUIWithTLS, self).setUp()
+        tls_cert, tls_key = self._create_tls_cert()
+        self.tls_cert_file = NamedTemporaryFile(delete=False)
+        self.tls_key_file = NamedTemporaryFile(delete=False)
+        with open(self.tls_cert_file.name, 'w') as f:
+            f.write(tls_cert.decode())
+        with open(self.tls_key_file.name, 'w') as f:
+            f.write(tls_key.decode())
+
+        parser = get_parser(default_config_files=[])
+        options = parser.parse_args([
+            "--tls-cert", self.tls_cert_file.name,
+            "--tls-key", self.tls_key_file.name,
+        ])
+        self.runner = Runner(self.environment)
+        self.stats = self.runner.stats
+        self.web_ui = self.environment.create_web_ui("127.0.0.1", 0, tls_cert=options.tls_cert, tls_key=options.tls_key)
+        gevent.sleep(0.01)
+        self.web_port = self.web_ui.server.server_port
+
+    def tearDown(self):
+        super(TestWebUIWithTLS, self).tearDown()
+        self.web_ui.stop()
+        self.runner.quit()
+        os.unlink(self.tls_cert_file.name)
+        os.unlink(self.tls_key_file.name)
+
+    def test_index_with_https(self):
+        # Suppress only the single warning from urllib3 needed.
+        from urllib3.exceptions import InsecureRequestWarning
+        requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)
+        self.assertEqual(200, requests.get("https://127.0.0.1:%i/" % self.web_port, verify=False).status_code)

locust/web.py

@@ -57,7 +57,7 @@ def my_custom_route():
     server = None
     """Reference to the :class:`pyqsgi.WSGIServer` instance"""
 
-    def __init__(self, environment, host, port, auth_credentials=None):
+    def __init__(self, environment, host, port, auth_credentials=None, tls_cert=None, tls_key=None):
         """
         Create WebUI instance and start running the web server in a separate greenlet (self.greenlet)
         
@@ -67,11 +67,15 @@ def __init__(self, environment, host, port, auth_credentials=None):
         port: Port that the web server should listen to
         auth_credentials:  If provided, it will enable basic auth with all the routes protected by default.
                            Should be supplied in the format: "user:pass".
+        tls_cert: A path to a TLS certificate
+        tls_key: A path to a TLS private key
         """
         environment.web_ui = self
         self.environment = environment
         self.host = host
         self.port = port
+        self.tls_cert = tls_cert
+        self.tls_key = tls_key
         app = Flask(__name__)
         self.app = app
         app.debug = True
@@ -276,7 +280,10 @@ def exceptions_csv():
         self.greenlet.link_exception(greenlet_exception_handler)
 
     def start(self):
-        self.server = pywsgi.WSGIServer((self.host, self.port), self.app, log=None)
+        if self.tls_cert and self.tls_key:
+            self.server = pywsgi.WSGIServer((self.host, self.port), self.app, log=None, keyfile=self.tls_key, certfile=self.tls_cert)
+        else:
+            self.server = pywsgi.WSGIServer((self.host, self.port), self.app, log=None)
         self.server.serve_forever()
 
     def stop(self):

setup.py

@@ -53,8 +53,9 @@
     ],
     test_suite="locust.test",
     tests_require=[
+        'cryptography',
         'mock',
-        'pyquery'
+        'pyquery',
         ],
     entry_points={
         'console_scripts': [

tox.ini

@@ -7,6 +7,7 @@ deps =
     flake8
     mock
     pyquery
+    cryptography
 commands =
     flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
     coverage run -m unittest discover []