Repository: locustio/locust
File: locust/web.py (Python)

Code before:
    from cgi import escape

from flask import Flask, make_response, jsonify, render_template, request
from gevent import pywsgi

from locust import __version__ as version
from io import StringIO

from . import runners
from .runners import MasterLocustRunner
from .stats import failures_csv, median_from_dict, requests_csv, sort_stats, stats_history_csv
from .util.cache import memoize
from .util.rounding import proper_round
from .util.timespan import parse_timespan

logger = logging.getLogger(__name__)

DEFAULT_CACHE_TIME = 2.0



class WebUI:
    server = None
    """Reference to pyqsgi.WSGIServer once it's started"""
    
    def __init__(self, environment, web_auth=None):
        environment.web_ui = self
        self.environment = environment
        app = Flask(__name__)
        self.app = app
        app.debug = True
        app.root_path = os.path.dirname(os.path.abspath(__file__))
        self.app.config["BASIC_AUTH_ENABLED"] = False
        self.auth = None

        if web_auth is not None:
            credentials = web_auth.split(':')
            if len(credentials) == 2:
                self.app.config["BASIC_AUTH_USERNAME"] = credentials[0]
                self.app.config["BASIC_AUTH_PASSWORD"] = credentials[1]
                self.app.config["BASIC_AUTH_ENABLED"] = True
                self.auth = BasicAuth()
                self.auth.init_app(self.app)
            else:
                sys.stderr.write("The credentials need to provided in the format username:password\n")
                sys.exit(1)


        def auth_required_if_enabled(view_func):
            @wraps(view_func)
            def wrapper(*args, **kwargs):

Review comment (security):
Maybe `auth_credentials` is a better variable name? `web_auth` is a good name for the CLI option, but in this context (constructor argument for the `WebUI` class) "web" feels redundant.

It would also be nice to have a short comment in the docstring of `__init__` that explains what the variable does and an example of a valid value.

Code after:
    from cgi import escape

from flask import Flask, make_response, jsonify, render_template, request
from gevent import pywsgi

from locust import __version__ as version
from io import StringIO

from . import runners
from .runners import MasterLocustRunner
from .stats import failures_csv, median_from_dict, requests_csv, sort_stats
from .util.cache import memoize
from .util.rounding import proper_round
from .util.timespan import parse_timespan

logger = logging.getLogger(__name__)

DEFAULT_CACHE_TIME = 2.0



class WebUI:
    server = None
    """Reference to pyqsgi.WSGIServer once it's started"""
    
    def __init__(self, environment, auth_credentials=None):
        """
        If auth_credentials is provided, it will enable basic auth with all the routes protected by default.
        Should be supplied in the format: "user:pass".
        """
        environment.web_ui = self
        self.environment = environment
        app = Flask(__name__)
        self.app = app
        app.debug = True
        app.root_path = os.path.dirname(os.path.abspath(__file__))
        self.app.config["BASIC_AUTH_ENABLED"] = False
        self.auth = None

        if auth_credentials is not None:
            credentials = auth_credentials.split(':')
            if len(credentials) == 2:
                self.app.config["BASIC_AUTH_USERNAME"] = credentials[0]
                self.app.config["BASIC_AUTH_PASSWORD"] = credentials[1]
                self.app.config["BASIC_AUTH_ENABLED"] = True
                self.auth = BasicAuth()
                self.auth.init_app(self.app)
            else:
                raise AuthCredentialsError("Invalid auth_credentials. It should be a string in the following format: 'user.pass'")

        @app.route('/')