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

Code before:

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):
                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

Review comment (security):
Instead of writing to stderr and calling `sys.exit()` we should raise an exception here. For example `InvalidCredentialsError` that could inherit from `ValueError`. Then we can catch this exception in main where we can output a proper error message with `logger.error()` and call `sys.exit(1)`.

Code after:

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('/')
        @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