Webui Add Custom Form to Auth Page

docs/extending-locust.rst

@@ -137,13 +137,14 @@ source code.
 Adding Authentication to the Web UI
 ===================================
 
-Locust uses `Flask-Login <https://pypi.org/project/Flask-Login/>`_ to handle authentication. The ``login_manager`` is
-exposed on ``environment.web_ui.app``, allowing the flexibility for you to implement any kind of auth that you would like!
+Locust uses `Flask-Login <https://pypi.org/project/Flask-Login/>`_ to handle authentication when the ``--web-login`` flag is present.
+The ``login_manager`` is exposed on ``environment.web_ui.app``, allowing the flexibility for you to implement any kind of auth that
+you would like!
 
 To use username / password authentication, simply provide a ``username_password_callback`` to the ``environment.web_ui.auth_args``.
 You are responsible for defining the route for the callback and implementing the authentication.
 
-Authentication providers can additionally be configured to allow authentication from 3rd parties such as GitHub or an SSO.
+Authentication providers can additionally be configured to allow authentication from 3rd parties such as GitHub or an SSO provider.
 Simply provide a list of desired ``auth_providers``. You may specify the ``label`` and ``icon`` for display on the button.
 The ``callback_url`` will be the url that the button directs to. You will be responsible for defining the callback route as
 well as the authentication with the 3rd party.
@@ -152,7 +153,17 @@ Whether you are using username / password authentication, an auth provider, or b
 to the ``login_manager``. The ``user_loader`` should return ``None`` to deny authentication or return a User object when
 authentication to the app should be granted.
 
-A full example can be seen `in the auth example <https://github.com/locustio/locust/tree/master/examples/web_ui_auth.py>`_.
+To display errors on the login page, such as an incorrect username / password combination, you may store the ``auth_error``
+on the session object: ``session["auth_error"] = "Incorrect username or password"``.
+
+A full example can be seen `in the auth example <https://github.com/locustio/locust/tree/master/examples/web_ui_auth/basic.py>`_.
+
+In certain situations you may wish to further extend the fields present in the auth form. To achieve this, pass a ``custom_form`` dict
+to the ``environment.web_ui.auth_args``. In this case, the fields will be represented by a list of ``inputs``, the callback url is
+configured by the ``custom_form.callback_url``, and the submit button may optionally be configured using the ``custom_form.submit_button_text``.
+The fields in the auth form may be a text, select, checkbox, or secret password field.
+
+For a full example see `configuring the custom_form in the auth example <https://github.com/locustio/locust/tree/master/examples/web_ui_auth/custom_form.py>`_.
 
 
 Run a background greenlet

examples/web_ui_auth.py → examples/web_ui_auth/basic.py

@@ -9,10 +9,9 @@
 
 from locust import HttpUser, events, task
 
-import json
 import os
 
-from flask import Blueprint, make_response, redirect, request, session, url_for
+from flask import Blueprint, redirect, request, session, url_for
 from flask_login import UserMixin, login_user
 
 
@@ -33,12 +32,12 @@ def get_id(self):
 auth_blueprint = Blueprint("auth", "web_ui_auth")
 
 
-def load_user(user_id):
-    return AuthUser(session.get("username"))
+def load_user(username):
+    return AuthUser(username)
 
 
 @events.init.add_listener
-def locust_init(environment, **kwargs):
+def locust_init(environment, **_kwargs):
     if environment.web_ui:
         environment.web_ui.login_manager.user_loader(load_user)
 
@@ -71,12 +70,11 @@ def login_submit():
 
             # Implement real password verification here
             if password:
-                session["username"] = username
                 login_user(AuthUser(username))
 
                 return redirect(url_for("index"))
 
-            environment.web_ui.auth_args = {**environment.web_ui.auth_args, "error": "Invalid username or password"}
+            session["auth_error"] = "Invalid username or password"
 
             return redirect(url_for("login"))
 

examples/web_ui_auth/custom_form.py

@@ -0,0 +1,107 @@
+"""
+Example of implementing authentication with a custom form for Locust when the --web-login
+flag is given
+
+This is only to serve as a starting point, proper authentication should be implemented
+according to your projects specifications.
+
+For more information, see https://docs.locust.io/en/stable/extending-locust.html#authentication
+"""
+
+from __future__ import annotations
+
+from locust import HttpUser, events, task
+
+import os
+
+from flask import Blueprint, redirect, request, session, url_for
+from flask_login import UserMixin, login_user
+
+
+class LocustHttpUser(HttpUser):
+    @task
+    def example(self):
+        self.client.get("/")
+
+
+class AuthUser(UserMixin):
+    def __init__(self, username):
+        self.username = username
+        self.is_admin = False
+        self.user_group: str | None = None
+
+    def get_id(self):
+        return self.username
+
+
+auth_blueprint = Blueprint("auth", "web_ui_auth")
+
+
+def load_user(user_id):
+    return AuthUser(user_id)
+
+
+@events.init.add_listener
+def locust_init(environment, **_kwargs):
+    if environment.web_ui:
+        environment.web_ui.login_manager.user_loader(load_user)
+
+        environment.web_ui.app.config["SECRET_KEY"] = os.getenv("FLASK_SECRET_KEY")
+
+        environment.web_ui.auth_args = {
+            "custom_form": {
+                "inputs": [
+                    {
+                        "label": "Username",
+                        "name": "username",
+                    },
+                    # boolean checkmark field
+                    {"label": "Admin", "name": "is_admin", "default_value": False},
+                    # select field
+                    {"label": "User Group", "name": "user_group", "choices": ["developer", "manager"]},
+                    {
+                        "label": "Password",
+                        "name": "password",
+                        "is_secret": True,
+                    },
+                    {
+                        "label": "Confirm Password",
+                        "name": "confirm_password",
+                        "is_secret": True,
+                    },
+                ],
+                "callback_url": "/login_submit",
+                "submit_button_text": "Submit",
+            },
+        }
+
+        @auth_blueprint.route("/login_submit", methods=["POST"])
+        def login_submit():
+            username = request.form.get("username")
+            password = request.form.get("password")
+            confirm_password = request.form.get("confirm_password")
+            is_admin = request.form.get("is_admin") == "on"
+            user_group = request.form.get("user_group")
+
+            if password != confirm_password:
+                session["auth_error"] = "Passwords do not match!"
+
+                return redirect(url_for("login"))
+
+            # Implement real password verification here
+            if password:
+                current_user = AuthUser(username)
+
+                # do something with your custom variables
+                current_user.is_admin = is_admin
+                current_user.user_group = user_group
+
+                login_user(AuthUser(username))
+
+                return redirect(url_for("index"))
+
+            session["auth_error"] = "Invalid username or password"
+
+            return redirect(url_for("login"))
+
+        environment.web_ui.app.register_blueprint(auth_blueprint)

locust/web.py

@@ -24,6 +24,7 @@
     request,
     send_file,
     send_from_directory,
+    session,
     url_for,
 )
 from flask_cors import CORS
@@ -504,6 +505,8 @@ def login():
             if not self.web_login:
                 return redirect(url_for("index"))
 
+            self.auth_args["error"] = session.get("auth_error", None)
+
             return render_template_from(
                 "auth.html",
                 auth_args=self.auth_args,
@@ -573,6 +576,7 @@ def my_custom_route():
         def wrapper(*args, **kwargs):
             if self.web_login:
                 try:
+                    session["auth_error"] = None
                     return login_required(view_func)(*args, **kwargs)
                 except Exception as e:
                     return f"Locust auth exception: {e} See https://docs.locust.io/en/stable/extending-locust.html#adding-authentication-to-the-web-ui for configuring authentication."

locust/webui/src/components/Form/CustomInput.tsx

@@ -0,0 +1,49 @@
+import { Checkbox, FormControlLabel, TextField } from '@mui/material';
+
+import PasswordField from 'components/Form/PasswordField';
+import Select from 'components/Form/Select';
+import { ICustomInput } from 'types/form.types';
+
+export default function CustomInput({
+  name,
+  label,
+  defaultValue,
+  choices,
+  isSecret,
+}: ICustomInput) {
+  if (choices) {
+    return (
+      <Select
+        defaultValue={defaultValue as string}
+        label={label}
+        name={name}
+        options={choices}
+        sx={{ width: '100%' }}
+      />
+    );
+  }
+
+  if (typeof defaultValue === 'boolean') {
+    return (
+      <FormControlLabel
+        control={<Checkbox defaultChecked={defaultValue} />}
+        label={label}
+        name={name}
+      />
+    );
+  }
+
+  if (isSecret) {
+    return <PasswordField defaultValue={defaultValue} label={label} name={name} />;
+  }
+
+  return (
+    <TextField
+      defaultValue={defaultValue}
+      label={label}
+      name={name}
+      sx={{ width: '100%' }}
+      type='text'
+    />
+  );
+}

locust/webui/src/components/Form/NumericField.tsx

@@ -20,7 +20,9 @@ export default function NumericField({ defaultValue, ...textFieldProps }: TextFi
   };
 
   useEffect(() => {
-    setValue(defaultValue as string);
+    if (defaultValue) {
+      setValue(defaultValue as string);
+    }
   }, [defaultValue]);
 
   return <TextField {...textFieldProps} onChange={filterNonNumeric} value={value} />;

locust/webui/src/components/Form/PasswordField.tsx

@@ -0,0 +1,35 @@
+import { useState } from 'react';
+import { Visibility, VisibilityOff } from '@mui/icons-material';
+import { FormControl, IconButton, InputAdornment, InputLabel, OutlinedInput } from '@mui/material';
+
+import { ICustomInput } from 'types/form.types';
+
+export default function PasswordField({
+  name = 'password',
+  label = 'Password',
+  defaultValue,
+}: Partial<Pick<ICustomInput, 'name' | 'label' | 'defaultValue'>>) {
+  const [showPassword, setShowPassword] = useState(false);
+
+  const handleClickShowPassword = () => setShowPassword(!showPassword);
+
+  return (
+    <FormControl variant='outlined'>
+      <InputLabel htmlFor={`${label}-${name}-field`}>{label}</InputLabel>
+      <OutlinedInput
+        defaultValue={defaultValue}
+        endAdornment={
+          <InputAdornment position='end'>
+            <IconButton edge='end' onClick={handleClickShowPassword}>
+              {showPassword ? <VisibilityOff /> : <Visibility />}
+            </IconButton>
+          </InputAdornment>
+        }
+        id={`${label}-${name}-field`}
+        label={label}
+        name={name}
+        type={showPassword ? 'text' : 'password'}
+      />
+    </FormControl>
+  );
+}