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

Code before:
    
    def __init__(self, environment):
        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__))
        app.config["BASIC_AUTH_ENABLED"] = False
        self.auth = BasicAuth()

        def auth_required_if_enabled(view_func):
            @wraps(view_func)
            def wrapper(*args, **kwargs):
                if app.config["BASIC_AUTH_ENABLED"]:
                    if self.auth.authenticate():
                        return view_func(*args, **kwargs)
                    else:
                        return self.auth.challenge()
                else:
                    return view_func(*args, **kwargs)
            return wrapper
        
        @app.route('/')
        @auth_required_if_enabled
        def index():
            if not environment.runner:
                return make_response("Error: Locust Environment does not have any runner", 500)
            
            is_distributed = isinstance(environment.runner, MasterLocustRunner)
            if is_distributed:
                worker_count = environment.runner.worker_count
            else:
                worker_count = 0
            
            override_host_warning = False
            if environment.host:
                host = environment.host
            elif environment.runner.locust_classes:
                all_hosts = set([l.host for l in environment.runner.locust_classes])
                if len(all_hosts) == 1:
                    host = list(all_hosts)[0]
                else:
                    # since we have mulitple Locust classes with different host attributes, we'll
                    # inform that specifying host will override the host for all Locust classes
                    override_host_warning = True
                    host = None
            else:
                host = None
            
            return render_template("index.html",

Review comment (security):
Looking at the [documentation](https://flask-basicauth.readthedocs.io/en/latest/) for Flask-BasicAuth, there seem to be an option called `BASIC_AUTH_FORCE`. Couldn't we use that and not have to decorate every endpoint?

Code after:
    """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('/')
        @self.auth_required_if_enabled
        def index():
            if not environment.runner:
                return make_response("Error: Locust Environment does not have any runner", 500)
            
            is_distributed = isinstance(environment.runner, MasterLocustRunner)
            if is_distributed:
                worker_count = environment.runner.worker_count
            else:
                worker_count = 0
            
            override_host_warning = False
            if environment.host:
                host = environment.host
            elif environment.runner.locust_classes:
                all_hosts = set([l.host for l in environment.runner.locust_classes])
                if len(all_hosts) == 1:
                    host = list(all_hosts)[0]
                else:
                    # since we have mulitple Locust classes with different host attributes, we'll
                    # inform that specifying host will override the host for all Locust classes
                    override_host_warning = True
                    host = None