locust-cloud 1.12.4__py3-none-any.whl → 1.13.0__py3-none-any.whl

locust_cloud/auth.py

@@ -1,443 +0,0 @@
-import logging
-import os
-from datetime import UTC, datetime, timedelta
-from typing import Any, TypedDict, cast
-
-import locust.env
-import requests
-import werkzeug
-from flask import Blueprint, redirect, request, session, url_for
-from flask_login import UserMixin, login_required, login_user, logout_user
-from locust.html import render_template_from
-from locust_cloud import __version__
-
-logger = logging.getLogger(__name__)
-
-
-class Credentials(TypedDict):
-    user_sub_id: str
-    refresh_token: str
-
-
-class AuthUser(UserMixin):
-    def __init__(self, user_sub_id: str):
-        self.user_sub_id = user_sub_id
-
-    def get_id(self):
-        return self.user_sub_id
-
-
-def set_credentials(username: str, credentials: Credentials, response: werkzeug.wrappers.response.Response):
-    if not credentials.get("user_sub_id"):
-        return response
-
-    user_sub_id = credentials["user_sub_id"]
-    refresh_token = credentials["refresh_token"]
-
-    response.set_cookie("username", username, expires=datetime.now(tz=UTC) + timedelta(days=365))
-    response.set_cookie("user_token", refresh_token, expires=datetime.now(tz=UTC) + timedelta(days=365))
-    response.set_cookie("user_sub_id", user_sub_id, expires=datetime.now(tz=UTC) + timedelta(days=365))
-
-    return response
-
-
-def register_auth(environment: locust.env.Environment):
-    environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "") + os.getenv("CUSTOMER_ID", "")
-    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")
-
-        if refresh_token:
-            environment.web_ui.template_args["username"] = username
-            return AuthUser(user_sub_id)
-
-        return None
-
-    environment.web_ui.login_manager.user_loader(load_user)
-    environment.web_ui.auth_args = cast(
-        Any,
-        {
-            "username_password_callback": f"{web_base_path}/authenticate",
-        },
-    )
-
-    environment.web_ui.auth_args["auth_providers"] = []
-    if environment.parsed_options.allow_signup:
-        environment.web_ui.auth_args["auth_providers"].append(
-            {"label": "Sign Up", "callback_url": f"{web_base_path}/signup"}
-        )
-    if environment.parsed_options.allow_forgot_password:
-        environment.web_ui.auth_args["auth_providers"].append(
-            {"label": "Forgot Password?", "callback_url": f"{web_base_path}/forgot-password"}
-        )
-
-    @auth_blueprint.route("/authenticate", methods=["POST"])
-    def login_submit():
-        username = request.form.get("username", "")
-        password = request.form.get("password")
-
-        try:
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/login",
-                json={"username": username, "password": password},
-                headers={"X-Client-Version": __version__},
-            )
-
-            auth_response.raise_for_status()
-
-            credentials = auth_response.json()
-
-            if credentials.get("challenge_session"):
-                session["challenge_session"] = credentials.get("challenge_session")
-                session["username"] = username
-
-                session["auth_error"] = ""
-
-                return redirect(url_for("locust_cloud_auth.password_reset"))
-            if os.getenv("CUSTOMER_ID", "") and credentials.get("customer_id") != os.getenv("CUSTOMER_ID", ""):
-                session["auth_error"] = "Invalid login for this deployment"
-                return redirect(url_for("locust.login"))
-
-            if not credentials.get("user_sub_id"):
-                session["auth_error"] = "Unknown error during authentication, check logs and/or contact support"
-                return redirect(url_for("locust.login"))
-
-            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:
-                session["auth_error"] = "Invalid username or password"
-            else:
-                logger.error(f"Unknown response from auth: {e.response.status_code} {e.response.text}")
-
-                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 environment.parsed_options.allow_signup:
-            return redirect(url_for("locust.login"))
-
-        if session.get("username"):
-            sign_up_args = {
-                "custom_form": {
-                    "inputs": [
-                        {
-                            "label": "Confirmation Code",
-                            "name": "confirmation_code",
-                            "is_required": True,
-                        },
-                    ],
-                    "callback_url": f"{web_base_path}/confirm-signup",
-                    "submit_button_text": "Confirm Email",
-                },
-            }
-        else:
-            sign_up_args = {
-                "custom_form": {
-                    "inputs": [
-                        {
-                            "label": "Username",
-                            "name": "username",
-                            "is_required": True,
-                            "type": "email",
-                        },
-                        {
-                            "label": "Full Name",
-                            "name": "customer_name",
-                            "is_required": True,
-                        },
-                        {
-                            "label": "Password",
-                            "name": "password",
-                            "is_secret": True,
-                            "is_required": True,
-                        },
-                        {
-                            "label": "Access Code",
-                            "name": "access_code",
-                            "is_required": True,
-                        },
-                        {
-                            "label": "I consent to:\n\n1. Only test your own website/service or our 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,
-                            "is_required": True,
-                        },
-                    ],
-                    "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 environment.parsed_options.allow_signup:
-            return redirect(url_for("locust.login"))
-
-        session["auth_sign_up_error"] = ""
-        session["auth_info"] = ""
-
-        username = request.form.get("username", "")
-        customer_name = request.form.get("customer_name", "")
-        password = request.form.get("password")
-        access_code = request.form.get("access_code")
-
-        try:
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/signup",
-                json={"username": username, "password": password, "access_code": access_code},
-            )
-
-            auth_response.raise_for_status()
-
-            session["user_sub_id"] = auth_response.json().get("user_sub_id")
-            session["username"] = username
-            session["customer_name"] = customer_name
-            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():
-        if not session.get("username"):
-            session["auth_sign_up_error"] = "An unexpected error occured. Please try again."
-            return redirect(url_for("locust_cloud_auth.signup"))
-
-        try:
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/resend-confirmation",
-                json={"username": session.get("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 environment.parsed_options.allow_signup:
-            return redirect(url_for("locust.login"))
-        if not session.get("user_sub_id"):
-            session["auth_sign_up_error"] = "An unexpected error occured. Please try again."
-            return redirect(url_for("locust_cloud_auth.signup"))
-
-        session["auth_sign_up_error"] = ""
-        confirmation_code = request.form.get("confirmation_code")
-
-        try:
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/confirm-signup",
-                json={
-                    "username": session.get("username"),
-                    "customer_name": session.get("customer_name"),
-                    "user_sub_id": session["user_sub_id"],
-                    "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("https://docs.locust.cloud/")
-        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("/forgot-password")
-    def forgot_password():
-        if not environment.parsed_options.allow_forgot_password:
-            return redirect(url_for("locust.login"))
-
-        forgot_password_args = {
-            "custom_form": {
-                "inputs": [
-                    {
-                        "label": "Username",
-                        "name": "username",
-                        "is_required": True,
-                        "type": "email",
-                    },
-                ],
-                "callback_url": f"{web_base_path}/send-forgot-password",
-                "submit_button_text": "Reset Password",
-            },
-            "info": "Enter your email and we will send a code to reset your password",
-        }
-
-        if session.get("auth_error"):
-            forgot_password_args["error"] = session["auth_error"]
-
-        return render_template_from("auth.html", auth_args=forgot_password_args)
-
-    @auth_blueprint.route("/send-forgot-password", methods=["POST"])
-    def send_forgot_password():
-        if not environment.parsed_options.allow_forgot_password:
-            return redirect(url_for("locust.login"))
-
-        try:
-            username = request.form.get("username", "")
-
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/forgot-password",
-                json={"username": username},
-            )
-
-            auth_response.raise_for_status()
-
-            session["username"] = username
-
-            return redirect(url_for("locust_cloud_auth.password_reset"))
-        except requests.exceptions.HTTPError as e:
-            message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
-            session["auth_info"] = ""
-            session["auth_error"] = message
-
-            return redirect(url_for("locust_cloud_auth.forgot_password"))
-
-    @auth_blueprint.route("/password-reset")
-    def password_reset():
-        if not environment.parsed_options.allow_forgot_password and not session.get("challenge_session"):
-            return redirect(url_for("locust.login"))
-
-        if session.get("challenge_session"):
-            reset_password_args = {
-                "custom_form": {
-                    "inputs": [
-                        {
-                            "label": "New Password",
-                            "name": "new_password",
-                            "is_required": True,
-                            "is_secret": True,
-                        },
-                    ],
-                    "callback_url": f"{web_base_path}/confirm-reset-password",
-                    "submit_button_text": "Set Password",
-                },
-                "info": "You must set a new password",
-            }
-        else:
-            reset_password_args = {
-                "custom_form": {
-                    "inputs": [
-                        {
-                            "label": "Confirmation Code",
-                            "name": "confirmation_code",
-                            "is_required": True,
-                        },
-                        {
-                            "label": "New Password",
-                            "name": "new_password",
-                            "is_required": True,
-                            "is_secret": True,
-                        },
-                    ],
-                    "callback_url": f"{web_base_path}/confirm-reset-password",
-                    "submit_button_text": "Reset Password",
-                },
-                "info": "Enter your the confirmation code that was sent to your email",
-            }
-
-        if session.get("auth_error"):
-            reset_password_args["error"] = session["auth_error"]
-
-        return render_template_from("auth.html", auth_args=reset_password_args)
-
-    @auth_blueprint.route("/confirm-reset-password", methods=["POST"])
-    def confirm_reset_password():
-        if not environment.parsed_options.allow_forgot_password and not session.get("challenge_session"):
-            return redirect(url_for("locust.login"))
-
-        try:
-            username = session["username"]
-            confirmation_code = request.form.get("confirmation_code")
-            new_password = request.form.get("new_password")
-
-            auth_response = requests.post(
-                f"{environment.parsed_options.deployer_url}/auth/password-reset",
-                json={
-                    "username": username,
-                    "confirmation_code": confirmation_code,
-                    "new_password": new_password,
-                    "challenge_session": session.get("challenge_session"),
-                },
-            )
-
-            auth_response.raise_for_status()
-
-            session["username"] = ""
-            session["auth_error"] = ""
-
-            if session.get("challenge_session"):
-                session["challenge_session"] = ""
-
-                return redirect(url_for("locust_cloud_auth.password_reset_success"))
-
-            session["auth_info"] = "Password reset successfully! Please login"
-
-            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_error"] = message
-
-            return redirect(url_for("locust_cloud_auth.password_reset"))
-
-    @auth_blueprint.route("/password-reset-success")
-    def password_reset_success():
-        return render_template_from(
-            "auth.html",
-            auth_args={
-                "info": "Password successfully set! Please review the [documentation](https://docs.locust.cloud/) and start your first testrun! If you have already ran some tests, you may also [login](/login)"
-            },
-        )
-
-    @auth_blueprint.route("/logout", methods=["POST"])
-    @login_required
-    def logout():
-        logout_user()
-        return redirect(url_for("locust.login"))
-
-    environment.web_ui.app.register_blueprint(auth_blueprint)

locust_cloud/credential_manager.py

@@ -1,141 +0,0 @@
-import logging
-import time
-from datetime import UTC, datetime
-from typing import Any
-
-import boto3
-import jwt
-import requests
-from botocore.credentials import RefreshableCredentials
-from botocore.session import Session as BotocoreSession
-from locust_cloud import __version__
-
-logger = logging.getLogger(__name__)
-
-
-class CredentialError(Exception):
-    """Custom exception for credential-related errors."""
-
-    pass
-
-
-class CredentialManager:
-    def __init__(
-        self,
-        lambda_url: str,
-        username: str | None = None,
-        password: str | None = None,
-        user_sub_id: str | None = None,
-        refresh_token: str | None = None,
-        access_key: str | None = None,
-        secret_key: str | None = None,
-    ) -> None:
-        self.lambda_url = lambda_url
-        self.username = username
-        self.password = password
-        self.user_sub_id = user_sub_id
-        self.refresh_token = refresh_token
-
-        self.credentials = {
-            "access_key": access_key,
-            "secret_key": secret_key,
-        }
-        self.cognito_client_id_token: str = ""
-        self.expiry_time: float = 0
-
-        self.obtain_credentials()
-
-        self.refreshable_credentials = RefreshableCredentials.create_from_metadata(
-            metadata=self.get_current_credentials(),
-            refresh_using=self.refresh_credentials,
-            method="custom-refresh",
-        )
-
-        botocore_session = BotocoreSession()
-        botocore_session._credentials = self.refreshable_credentials  # type: ignore
-        botocore_session.set_config_variable("signature_version", "v4")
-
-        self.session = boto3.Session(botocore_session=botocore_session)
-        logger.debug("Boto3 session created with RefreshableCredentials.")
-
-    def obtain_credentials(self) -> None:
-        payload = {}
-        if self.username and self.password:
-            payload = {"username": self.username, "password": self.password}
-        elif self.user_sub_id and self.refresh_token:
-            payload = {"user_sub_id": self.user_sub_id, "refresh_token": self.refresh_token}
-        else:
-            raise CredentialError("Insufficient credentials to obtain AWS session.")
-
-        try:
-            response = requests.post(
-                f"{self.lambda_url}/auth/login",
-                json=payload,
-                headers={"X-Client-Version": __version__},
-            )
-            response.raise_for_status()
-            data = response.json()
-
-            token_key = next(
-                (key for key in ["cognito_client_id_token", "id_token", "access_token"] if key in data), None
-            )
-
-            if not token_key:
-                raise CredentialError("No valid token found in authentication response.")
-
-            self.credentials = {
-                "access_key": data.get("aws_access_key_id"),
-                "secret_key": data.get("aws_secret_access_key"),
-                "token": data.get("aws_session_token"),
-            }
-
-            token = data.get(token_key)
-            if not token:
-                raise CredentialError(f"Token '{token_key}' is missing in the authentication response.")
-
-            decoded = jwt.decode(token, options={"verify_signature": False})
-            self.expiry_time = decoded.get("exp", time.time() + 3600) - 60  # Refresh 1 minute before expiry
-
-            self.cognito_client_id_token = token
-
-        except requests.exceptions.HTTPError as http_err:
-            response = http_err.response
-            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}"
-                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
-        except jwt.DecodeError as decode_err:
-            raise CredentialError(f"Failed to decode JWT token: {decode_err}") from decode_err
-        except KeyError as key_err:
-            raise CredentialError(f"Missing expected key in authentication response: {key_err}") from key_err
-
-    def refresh_credentials(self) -> dict[str, Any]:
-        logger.debug("Refreshing credentials using refresh_credentials method.")
-        self.obtain_credentials()
-        return {
-            "access_key": self.credentials.get("access_key"),
-            "secret_key": self.credentials.get("secret_key"),
-            "token": self.credentials.get("token"),
-            "expiry_time": datetime.fromtimestamp(self.expiry_time, tz=UTC).isoformat(),
-        }
-
-    def get_current_credentials(self) -> dict[str, Any]:
-        if not self.cognito_client_id_token:
-            raise CredentialError("cognito_client_id_token not set in CredentialManager.")
-
-        return {
-            "access_key": self.credentials.get("access_key"),
-            "secret_key": self.credentials.get("secret_key"),
-            "token": self.credentials.get("token"),
-            "expiry_time": datetime.fromtimestamp(self.expiry_time, tz=UTC).isoformat(),
-            "cognito_client_id_token": self.cognito_client_id_token,
-        }

locust_cloud/idle_exit.py

@@ -1,38 +0,0 @@
-import logging
-import sys
-
-import gevent
-import locust.env
-from locust import events
-
-logger = logging.getLogger(__name__)
-
-
-class IdleExit:
-    def __init__(self, environment: locust.env.Environment):
-        self.environment = environment
-        self._destroy_task: gevent.Greenlet | None = None
-        events.test_start.add_listener(self.on_locust_state_change)
-        events.test_stop.add_listener(self.on_test_stop)
-        events.quit.add_listener(self.on_locust_state_change)
-
-        if not self.environment.parsed_options.autostart:
-            self._destroy_task = gevent.spawn(self._destroy)
-
-    def _destroy(self):
-        gevent.sleep(1800)
-        logger.info("Locust was detected as idle (no test running) for more than 30 minutes")
-        self.environment.runner.quit()
-
-        if self.environment.web_ui:
-            self.environment.web_ui.greenlet.kill(timeout=5)
-
-            if self.environment.web_ui.greenlet.started:
-                sys.exit(1)
-
-    def on_test_stop(self, **kwargs):
-        self._destroy_task = gevent.spawn(self._destroy)
-
-    def on_locust_state_change(self, **kwargs):
-        if self._destroy_task:
-            self._destroy_task.kill()