locust-cloud 1.8.0__py3-none-any.whl → 1.9.0__py3-none-any.whl

locust_cloud/auth.py

@@ -1,17 +1,20 @@
 import logging
 import os
 from datetime import UTC, datetime, timedelta
-from typing import TypedDict
+from typing import Any, TypedDict, cast
 
 import locust.env
 import requests
 import werkzeug
-from flask import redirect, request, url_for
+from flask import Blueprint, redirect, request, session, url_for
 from flask_login import UserMixin, login_user
+from locust.html import render_template_from
 from locust_cloud import __version__
-from locust_cloud.constants import DEFAULT_DEPLOYER_URL
 
-DEPLOYER_URL = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", DEFAULT_DEPLOYER_URL)
+REGION = os.environ.get("AWS_DEFAULT_REGION")
+DEPLOYER_URL = f"https://api.{REGION}.locust.cloud/1"
+ALLOW_SIGNUP = os.environ.get("ALLOW_SIGNUP", True)
+
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +50,9 @@ def register_auth(environment: locust.env.Environment):
     environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
     environment.web_ui.app.debug = False
 
+    web_base_path = environment.parsed_options.web_base_path
+    auth_blueprint = Blueprint("locust_cloud_auth", __name__, url_prefix=web_base_path)
+
     def load_user(user_sub_id: str):
         username = request.cookies.get("username")
         refresh_token = request.cookies.get("user_token")
@@ -58,11 +64,19 @@ def register_auth(environment: locust.env.Environment):
         return None
 
     environment.web_ui.login_manager.user_loader(load_user)
-    environment.web_ui.auth_args = {
-        "username_password_callback": "/authenticate",
-    }
+    environment.web_ui.auth_args = cast(
+        Any,
+        {
+            "username_password_callback": f"{web_base_path}/authenticate",
+        },
+    )
+
+    if ALLOW_SIGNUP:
+        environment.web_ui.auth_args["auth_providers"] = [
+            {"label": "Sign Up", "callback_url": f"{web_base_path}/signup"}
+        ]
 
-    @environment.web_ui.app.route("/authenticate", methods=["POST"])
+    @auth_blueprint.route("/authenticate", methods=["POST"])
     def login_submit():
         username = request.form.get("username", "")
         password = request.form.get("password")
@@ -77,19 +91,168 @@ def register_auth(environment: locust.env.Environment):
             auth_response.raise_for_status()
 
             credentials = auth_response.json()
-            response = redirect(url_for("index"))
+            response = redirect(url_for("locust.index"))
             response = set_credentials(username, credentials, response)
             login_user(AuthUser(credentials["user_sub_id"]))
 
             return response
         except requests.exceptions.HTTPError as e:
             if e.response.status_code == 401:
-                environment.web_ui.auth_args["error"] = "Invalid username or password"
+                session["auth_error"] = "Invalid username or password"
             else:
                 logger.error(f"Unknown response from auth: {e.response.status_code} {e.response.text}")
 
-                environment.web_ui.auth_args["error"] = (
-                    "Unknown error during authentication, check logs and/or contact support"
-                )
+                session["auth_error"] = "Unknown error during authentication, check logs and/or contact support"
+
+            return redirect(url_for("locust.login"))
+
+    @auth_blueprint.route("/signup")
+    def signup():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        if session.get("username"):
+            sign_up_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "Confirmation Code",
+                            "name": "confirmation_code",
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/confirm-signup",
+                    "submit_button_text": "Confirm Email",
+                },
+            }
+        else:
+            sign_up_args = {
+                "custom_form": {
+                    "inputs": [
+                        {
+                            "label": "Username",
+                            "name": "username",
+                        },
+                        {
+                            "label": "Password",
+                            "name": "password",
+                            "is_secret": True,
+                        },
+                        {
+                            "label": "Access Code",
+                            "name": "access_code",
+                        },
+                        {
+                            "label": "I consent to:\n\n1. Only test your own website/service or our example example target\n\n2. Only use locust-cloud for its intended purpose: to load test other sites/services.\n\n3. Not attempt to circumvent your account limitations (e.g. max user count or max request count)\n\n4. Not use personal data (real names, addresses etc) in your tests.",
+                            "name": "consent",
+                            "default_value": False,
+                        },
+                    ],
+                    "callback_url": f"{web_base_path}/create-account",
+                    "submit_button_text": "Sign Up",
+                },
+            }
+
+        if session.get("auth_info"):
+            sign_up_args["info"] = session["auth_info"]
+        if session.get("auth_sign_up_error"):
+            sign_up_args["error"] = session["auth_sign_up_error"]
+
+        return render_template_from(
+            "auth.html",
+            auth_args=sign_up_args,
+        )
+
+    @auth_blueprint.route("/create-account", methods=["POST"])
+    def create_account():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        session["auth_sign_up_error"] = ""
+        session["auth_info"] = ""
+
+        username = request.form.get("username", "")
+        password = request.form.get("password")
+        access_code = request.form.get("access_code")
+        has_consented = request.form.get("consent")
+
+        if not has_consented:
+            session["auth_sign_up_error"] = "Please accept the terms and conditions to create an account."
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+        try:
+            auth_response = requests.post(
+                f"{DEPLOYER_URL}/auth/signup",
+                json={"username": username, "password": password, "access_code": access_code},
+            )
+
+            auth_response.raise_for_status()
+
+            session["user_sub"] = auth_response.json().get("user_sub")
+            session["username"] = username
+            session["auth_info"] = (
+                "Please check your email and enter the confirmation code. If you didn't get a code after one minute, you can [request a new one](/resend-code)"
+            )
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+    @auth_blueprint.route("/resend-code")
+    def resend_code():
+        try:
+            auth_response = requests.post(
+                "http://localhost:8000/1/auth/resend-confirmation",
+                json={"username": session["username"]},
+            )
+
+            auth_response.raise_for_status()
+
+            session["auth_sign_up_error"] = ""
+            session["auth_info"] = "Confirmation code sent, please check your email."
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
+
+    @auth_blueprint.route("/confirm-signup", methods=["POST"])
+    def confirm_signup():
+        if not ALLOW_SIGNUP:
+            return redirect(url_for("locust.login"))
+
+        session["auth_sign_up_error"] = ""
+        confirmation_code = request.form.get("confirmation_code")
+
+        try:
+            auth_response = requests.post(
+                f"{DEPLOYER_URL}/auth/confirm-signup",
+                json={
+                    "username": session.get("username"),
+                    "user_sub": session["user_sub"],
+                    "confirmation_code": confirmation_code,
+                },
+            )
+
+            auth_response.raise_for_status()
+
+            session["username"] = None
+            session["auth_info"] = "Account created successfully!"
+            session["auth_sign_up_error"] = ""
+
+            return redirect(url_for("locust.login"))
+        except requests.exceptions.HTTPError as e:
+            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
+            session["auth_info"] = ""
+            session["auth_sign_up_error"] = message
+
+            return redirect(url_for("locust_cloud_auth.signup"))
 
-            return redirect(url_for("login"))
+    environment.web_ui.app.register_blueprint(auth_blueprint)

locust_cloud/cloud.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import math
 import os
 import sys
 import time
@@ -14,13 +13,6 @@ import configargparse
 import requests
 from botocore.exceptions import ClientError
 from locust_cloud import __version__
-from locust_cloud.constants import (
-    DEFAULT_CLUSTER_NAME,
-    DEFAULT_DEPLOYER_URL,
-    DEFAULT_NAMESPACE,
-    DEFAULT_REGION_NAME,
-    USERS_PER_WORKER,
-)
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
 
@@ -110,29 +102,9 @@ advanced.add_argument(
 advanced.add_argument(
     "--region",
     type=str,
-    default=os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION_NAME),
+    default=os.environ.get("AWS_DEFAULT_REGION"),
     help="Sets the AWS region to use for the deployed cluster, e.g. us-east-1. It defaults to use AWS_DEFAULT_REGION env var, like AWS tools.",
 )
-advanced.add_argument(
-    "--kube-cluster-name",
-    type=str,
-    default=DEFAULT_CLUSTER_NAME,
-    help=configargparse.SUPPRESS,
-    env_var="KUBE_CLUSTER_NAME",
-)
-advanced.add_argument(
-    "--kube-namespace",
-    type=str,
-    default=DEFAULT_NAMESPACE,
-    help=configargparse.SUPPRESS,
-    env_var="KUBE_NAMESPACE",
-)
-parser.add_argument(
-    "--deployer-url",
-    type=str,
-    default=DEFAULT_DEPLOYER_URL,
-    help=configargparse.SUPPRESS,
-)
 parser.add_argument(
     "--aws-access-key-id",
     type=str,
@@ -162,7 +134,7 @@ parser.add_argument(
 parser.add_argument(
     "--workers",
     type=int,
-    help=f"Number of workers to use for the deployment. Defaults to number of users divided by {USERS_PER_WORKER}.",
+    help="Number of workers to use for the deployment. Defaults to number of users divided by 500, but the default may be customized for your account.",
     default=None,
 )
 parser.add_argument(
@@ -194,25 +166,30 @@ logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
 
+api_url = f"https://api.{options.region}.locust.cloud/1"
+
+
 def main() -> None:
-    s3_bucket = f"{options.kube_cluster_name}-{options.kube_namespace}"
+    if not options.region:
+        logger.error(
+            "Setting a region is required to use Locust Cloud. Please ensure the AWS_DEFAULT_REGION env variable or the --region flag is set."
+        )
+        sys.exit(1)
+
+    s3_bucket = "dmdb-default" if options.region == "us-east-1" else "locust-default"
     deployments: list[Any] = []
-    worker_count: int = max(options.workers or math.ceil(options.users / USERS_PER_WORKER), 2)
-    os.environ["AWS_DEFAULT_REGION"] = options.region
 
-    try:
-        if not (
-            (options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)
-        ):
-            logger.error(
-                "Authentication is required to use Locust Cloud. Please ensure the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
-            )
-            sys.exit(1)
+    if not ((options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)):
+        logger.error(
+            "Authentication is required to use Locust Cloud. Please ensure the LOCUST_CLOUD_USERNAME and LOCUST_CLOUD_PASSWORD environment variables are set."
+        )
+        sys.exit(1)
 
-        logger.info(f"Authenticating ({os.environ['AWS_DEFAULT_REGION']}, v{__version__})")
-        logger.debug(f"Lambda url: {options.deployer_url}")
+    try:
+        logger.info(f"Authenticating ({options.region}, v{__version__})")
+        logger.debug(f"Lambda url: {api_url}")
         credential_manager = CredentialManager(
-            lambda_url=options.deployer_url,
+            lambda_url=api_url,
             access_key=options.aws_access_key_id,
             secret_key=options.aws_secret_access_key,
             username=options.username,
@@ -279,20 +256,21 @@ def main() -> None:
             ]
             and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{options.deployer_url}/{options.kube_cluster_name}"
+        deploy_endpoint = f"{api_url}/deploy"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_LOCUSTFILE", "value": locustfile_url},
                 {"name": "LOCUST_USERS", "value": str(options.users)},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
                 {"name": "LOCUSTCLOUD_REQUIREMENTS_URL", "value": requirements_url},
-                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": options.deployer_url},
+                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": api_url},
                 *locust_env_variables,
             ],
-            "worker_count": worker_count,
             "user_count": options.users,
             "image_tag": options.image_tag,
         }
+        if options.workers is not None:
+            payload["worker_count"] = options.workers
         headers = {
             "Authorization": f"Bearer {cognito_client_id_token}",
             "Content-Type": "application/json",
@@ -325,7 +303,7 @@ def main() -> None:
         logger.debug("Interrupted by user")
         sys.exit(0)
 
-    log_group_name = f"/eks/{options.kube_cluster_name}-{options.kube_namespace}"
+    log_group_name = "/eks/dmdb-default" if options.region == "us-east-1" else "/eks/locust-default"
     master_pod_name = next((deployment for deployment in deployments if "master" in deployment), None)
 
     if not master_pod_name:
@@ -420,9 +398,8 @@ def delete(s3_bucket, credential_manager):
             headers["AWS_SESSION_TOKEN"] = token
 
         response = requests.delete(
-            f"{options.deployer_url}/{options.kube_cluster_name}",
+            f"{api_url}/teardown",
             headers=headers,
-            params={"namespace": options.kube_namespace} if options.kube_namespace else {},
         )
 
         if response.status_code != 200:

locust_cloud/credential_manager.py

@@ -100,13 +100,16 @@ class CredentialManager:
 
         except requests.exceptions.HTTPError as http_err:
             response = http_err.response
-            if response is not None and response.status_code == 401:
+            if response is None:
+                raise CredentialError("Response was None?!") from http_err
+
+            if response.status_code == 401:
                 raise CredentialError("Incorrect username or password.") from http_err
             else:
                 if js := response.json():
                     if message := js.get("Message"):
                         raise CredentialError(message)
-                error_info = f"HTTP {response.status_code} {response.reason}" if response else "No response received."
+                error_info = f"HTTP {response.status_code} {response.reason}"
                 raise CredentialError(f"HTTP error occurred while obtaining credentials: {error_info}") from http_err
         except requests.exceptions.RequestException as req_err:
             raise CredentialError(f"Request exception occurred while obtaining credentials: {req_err}") from req_err

locust_cloud/timescale/query.py

@@ -1,6 +1,6 @@
 import logging
 
-from flask import make_response, request
+from flask import Blueprint, make_response, request
 from flask_login import login_required
 from locust_cloud.timescale.queries import queries
 
@@ -12,7 +12,11 @@ def adapt_timestamp(result):
 
 
 def register_query(environment, pool):
-    @environment.web_ui.app.route("/cloud-stats/<query>", methods=["POST"])
+    cloud_stats_blueprint = Blueprint(
+        "locust_cloud_stats", __name__, url_prefix=environment.parsed_options.web_base_path
+    )
+
+    @cloud_stats_blueprint.route("/cloud-stats/<query>", methods=["POST"])
     @login_required
     def query(query):
         results = []
@@ -57,7 +61,9 @@ def register_query(environment, pool):
                 return results
             else:
                 logger.warning(f"Received invalid query key: '{query}'")
-                return make_response("Invalid query key", 401)
+                return make_response({"error": "Invalid query key"}, 401)
         except Exception as e:
             logger.info(f"Error executing UI query '{query}': {e}", exc_info=True)
-            return make_response("Error executing query", 401)
+            return make_response({"error": "Error executing query"}, 401)
+
+    environment.web_ui.app.register_blueprint(cloud_stats_blueprint)

locust_cloud/webui/.eslintrc

@@ -29,7 +29,7 @@
             "group": "internal",
           },
           {
-            "pattern": "{components,hooks,redux,utils}/**",
+            "pattern": "{components,hooks,redux,types,utils}/**",
             "group": "internal",
           },
         ],