Azure AD Authentication and Testing

from locust import HttpUser, task, between
import requests
import json

class PortalUser(HttpUser):
    wait_time = between(1, 3)  # Simulated user wait time
    token_url = "https://login.microsoftonline.com/YOUR_TENANT_ID/oauth2/v2.0/token"
    client_id = "YOUR_CLIENT_ID"
    client_secret = "YOUR_CLIENT_SECRET"
    scope = "https://graph.microsoft.com/.default"
    token = None
    headers = {}
    cookies = {}

    def on_start(self):
        """Fetch OAuth2 Token and store it for use in all requests."""
        self.token = self.get_oauth_token()
        self.headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

        # Initial request to fetch session cookies
        self.login_portal()

    def get_oauth_token(self):
        """Retrieve OAuth token from Azure AD."""
        data = {
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "scope": self.scope,
            "grant_type": "client_credentials",
        }
        response = requests.post(self.token_url, data=data)
        response.raise_for_status()  # Ensure request succeeded
        return response.json().get("access_token")

    def login_portal(self):
        """Initial request to portal to retrieve required cookies."""
        response = self.client.get("/login", headers=self.headers)
        self.cookies.update(response.cookies.get_dict())  # Store cookies

    @task
    def perform_task(self):
        """Perform a simulated user action with correct authentication headers."""
        self.headers["X-CSRF-Token"] = self.cookies.get("AspNetCore.Antiforgery", "")

        response = self.client.get(
            "/dashboard",
            headers=self.headers,
            cookies=self.cookies
        )

        if "AspNetCore.Antiforgery" in response.cookies:
            self.cookies["AspNetCore.Antiforgery"] = response.cookies["AspNetCore.Antiforgery"]

        # Simulated POST action (e.g., event signup)
        self.client.post(
            "/event/signup",
            headers=self.headers,
            cookies=self.cookies,
            json={"event_id": 123}
        )