From e59a10278082907f02e36364bba71fcae819c5dc Mon Sep 17 00:00:00 2001 From: Matt Dodge Date: Sat, 25 Apr 2020 15:59:32 -0700 Subject: [PATCH 1/4] Add CLI params for TLS cert and key - serves over HTTPS --- locust/argument_parser.py | 10 ++++++ locust/env.py | 4 +-- locust/main.py | 6 ++-- locust/test/test_web.py | 71 +++++++++++++++++++++++++++++++++++++++ locust/web.py | 11 ++++-- 5 files changed, 96 insertions(+), 6 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 59ceb419f1..4c36396f11 100644 --- a/locust/argument_parser.py +++ b/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", diff --git a/locust/env.py b/locust/env.py index d7b06cd671..58db73510a 100644 --- a/locust/env.py +++ b/locust/env.py @@ -116,7 +116,7 @@ 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 ` instance for this Environment and start running the web server @@ -125,5 +125,5 @@ def create_web_ui(self, host="", port=8089, auth_credentials=None): :param port: Port that the web server should listen to :param auth_credentials: If provided (in format "username:password") basic auth will be enabled """ - 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 diff --git a/locust/main.py b/locust/main.py index dcf7be8dfc..44f789c431 100644 --- a/locust/main.py +++ b/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) diff --git a/locust/test/test_web.py b/locust/test/test_web.py index d05ba97ce5..47e4cd3a8d 100644 --- a/locust/test/test_web.py +++ b/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) diff --git a/locust/web.py b/locust/web.py index 4e4c4e7399..b033c0f4c4 100644 --- a/locust/web.py +++ b/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): From e7d2a0058f4c87b6b3e9892a258c46ef2ec5e948 Mon Sep 17 00:00:00 2001 From: Matt Dodge Date: Sat, 25 Apr 2020 16:08:20 -0700 Subject: [PATCH 2/4] Inclue cryptography in setup.py test reqs --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index b107b81301..63d93509cd 100644 --- a/setup.py +++ b/setup.py @@ -53,8 +53,9 @@ ], test_suite="locust.test", tests_require=[ + 'cryptography', 'mock', - 'pyquery' + 'pyquery', ], entry_points={ 'console_scripts': [ From 15945757c180fff9b710c2814f20a9cf9f8e8da8 Mon Sep 17 00:00:00 2001 From: Matt Dodge Date: Sat, 25 Apr 2020 16:12:29 -0700 Subject: [PATCH 3/4] Add cryptography to tox test deps too --- tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/tox.ini b/tox.ini index dbf9a4c4bc..4f4ee24daa 100644 --- a/tox.ini +++ b/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 [] From eb19c57c508b0f63221821635b867038739d3423 Mon Sep 17 00:00:00 2001 From: Matt Dodge Date: Mon, 27 Apr 2020 08:26:26 -0700 Subject: [PATCH 4/4] Add docstring for tls cert and key --- locust/env.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/locust/env.py b/locust/env.py index 58db73510a..56ba2c3604 100644 --- a/locust/env.py +++ b/locust/env.py @@ -124,6 +124,10 @@ def create_web_ui(self, host="", port=8089, auth_credentials=None, tls_cert=None 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, tls_cert=tls_cert, tls_key=tls_key) return self.web_ui