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

Code before:
    if options.run_time:
        if not options.no_web:
            logger.error("The --run-time argument can only be used together with --no-web")
            sys.exit(1)
        if options.worker:
            logger.error("--run-time should be specified on the master node, and not on worker nodes")
            sys.exit(1)
        try:
            options.run_time = parse_timespan(options.run_time)
        except ValueError:
            logger.error("Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
            sys.exit(1)
        def spawn_run_time_limit_greenlet():
            logger.info("Run time limit set to %s seconds" % options.run_time)
            def timelimit_stop():
                logger.info("Time limit reached. Stopping Locust.")
                runner.quit()
            gevent.spawn_later(options.run_time, timelimit_stop)
    
    # start Web UI
    if not options.no_web and not options.worker:
        # spawn web greenlet
        logger.info("Starting web monitor at http://%s:%s" % (options.web_host or "*", options.web_port))
        web_ui = WebUI(environment=environment)
        if options.web_auth is not None:
            web_ui.app.config["BASIC_AUTH_ENABLED"] = True
            credentials = options.web_auth.split(':')
            if len(credentials) == 2:
                web_ui.app.config["BASIC_AUTH_USERNAME"] = credentials[0]
                web_ui.app.config["BASIC_AUTH_PASSWORD"] = credentials[1]
                web_ui.auth.init_app(web_ui.app)
            else:
                sys.stderr.write("The credentials need to provided in the format username:password\n")
                sys.exit(1)
        main_greenlet = gevent.spawn(web_ui.start, host=options.web_host, port=options.web_port)
    else:
        web_ui = None
    
    # Fire locust init event which can be used by end-users' code to run setup code that
    # need access to the Environment, Runner or WebUI
    environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui)    
    
    if options.no_web:
        # headless mode
        if options.master:
            # what for worker nodes to connect
            while len(runner.clients.ready) < options.expect_workers:
                logging.info("Waiting for workers to be ready, %s of %s connected",
                             len(runner.clients.ready), options.expect_workers)
                time.sleep(1)
        if not options.worker:

Review comment (security):
I think it would be better to add an `auth` argument to `WebUI.__init__` (that defaults to None) and move most of this logic to the WebUI class.

Code after:
    if options.run_time:
        if not options.headless:
            logger.error("The --run-time argument can only be used together with --headless")
            sys.exit(1)
        if options.worker:
            logger.error("--run-time should be specified on the master node, and not on worker nodes")
            sys.exit(1)
        try:
            options.run_time = parse_timespan(options.run_time)
        except ValueError:
            logger.error("Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
            sys.exit(1)
        def spawn_run_time_limit_greenlet():
            logger.info("Run time limit set to %s seconds" % options.run_time)
            def timelimit_stop():
                logger.info("Time limit reached. Stopping Locust.")
                runner.quit()
            gevent.spawn_later(options.run_time, 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 or "*", options.web_port))
        try:
            web_ui = WebUI(environment=environment, auth_credentials=options.web_auth)
        except AuthCredentialsError:
            logger.error("Credentials supplied with --web-auth should have the format: username:password")
            sys.exit(1)
        else:
            main_greenlet = gevent.spawn(web_ui.start, host=options.web_host, port=options.web_port)
    else:
        web_ui = None
    
    # Fire locust init event which can be used by end-users' code to run setup code that
    # need access to the Environment, Runner or WebUI
    environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui)    
    
    if options.headless:
        # headless mode
        if options.master:
            # what for worker nodes to connect
            while len(runner.clients.ready) < options.expect_workers:
                logging.info("Waiting for workers to be ready, %s of %s connected",
                             len(runner.clients.ready), options.expect_workers)
                time.sleep(1)
        if not options.worker:
            # start the test
            if options.step_time:
                runner.start_stepload(options.num_clients, options.hatch_rate, options.step_clients, options.step_time)
            else:
                runner.start(options.num_clients, options.hatch_rate)