1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
|
"""
Example of implementing authentication 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 locust import HttpUser, events, task
import json
import os
from flask import Blueprint, make_response, 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
def get_id(self):
return self.username
auth_blueprint = Blueprint("auth", "web_ui_auth")
def load_user(user_id):
return AuthUser(session.get("username"))
@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 = {
"username_password_callback": "/login_submit",
"auth_providers": [
{
"label": "Github",
"callback_url": "/login/github",
"icon_url": "https://static-00.iconduck.com/assets.00/github-icon-1024x994-4h5sdmko.png",
},
],
}
@auth_blueprint.route("/login/github")
def google_login():
# Implement authentication with desired auth provider
username = "username"
session["username"] = username
login_user(AuthUser("username"))
return redirect(url_for("index"))
@auth_blueprint.route("/login_submit")
def login_submit():
username = request.args.get("username")
password = request.args.get("password")
# 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"}
return redirect(url_for("login"))
environment.web_ui.app.register_blueprint(auth_blueprint)
|
