"""Interactive OAuth authentication service for the dev server. Implements the same PKCE-based authorization code flow as the Python SDK's ``uipath auth`` CLI command. A temporary callback server is spun up on one of the registered redirect-URI ports (8104, 8055, 42042) to receive the token from the browser. """ from __future__ import annotations import asyncio import base64 import hashlib import http.server import json import logging import os import socketserver import threading import time from dataclasses import dataclass, field from pathlib import Path from typing import Any from urllib.parse import urlencode import httpx logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- CLIENT_ID = "36dea5b8-e8bb-423d-8e7b-c808df8f1c00" SCOPES = ( "offline_access ProcessMining OrchestratorApiUserAccess StudioWebBackend " "IdentityServerApi ConnectionService DataService DocumentUnderstanding " "Du.Digitization.Api Du.Classification.Api Du.Extraction.Api " "Du.Validation.Api EnterpriseContextService Directory JamJamApi " "LLMGateway LLMOps OMS RCS.FolderAuthorization TM.Projects " "TM.TestCases TM.Requirements TM.TestSets AutomationSolutions" ) CANDIDATE_PORTS = [8104, 8055, 42042] # --------------------------------------------------------------------------- # PKCE helpers # --------------------------------------------------------------------------- def _generate_pkce() -> tuple[str, str]: """Return (code_verifier, code_challenge) for PKCE S256.""" verifier = base64.urlsafe_b64encode(os.urandom(32)).decode().rstrip("=") challenge = ( base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) .decode() .rstrip("=") ) return verifier, challenge def _generate_state() -> str: return base64.urlsafe_b64encode(os.urandom(32)).decode().rstrip("=") def _parse_jwt_payload(token: str) -> dict[str, Any]: """Decode the payload section of a JWT (no signature verification).""" parts = token.split(".") if len(parts) < 2: raise ValueError("Invalid JWT") padded = parts[1] + "=" * (-len(parts[1]) % 4) return json.loads(base64.urlsafe_b64decode(padded)) # --------------------------------------------------------------------------- # Port finder # --------------------------------------------------------------------------- def _try_bind_server( candidates: list[int], handler_class: type ) -> socketserver.TCPServer | None: """Try to bind a TCPServer on the first available candidate port.""" socketserver.TCPServer.allow_reuse_address = True for port in candidates: try: httpd = socketserver.TCPServer(("127.0.0.1", port), handler_class) return httpd except OSError: continue return None # --------------------------------------------------------------------------- # Auth state # --------------------------------------------------------------------------- @dataclass class AuthState: """Mutable state for the in-progress or completed OAuth flow.""" status: str = "unauthenticated" # unauthenticated | pending | needs_tenant | authenticated | expired environment: str = "cloud" token_data: dict[str, Any] = field(default_factory=dict) tenants: list[dict[str, str]] = field(default_factory=list) organization: dict[str, str] = field(default_factory=dict) uipath_url: str | None = None # Remembered from last session for seamless re-auth _last_tenant: str | None = None _last_org: dict[str, str] = field(default_factory=dict) _last_environment: str | None = None # internal _code_verifier: str | None = None _state: str | None = None _port: int | None = None _callback_server: _CallbackServer | None = None _token_event: asyncio.Event | None = None _loop: asyncio.AbstractEventLoop | None = None _wait_task: asyncio.Task[None] | None = None _auth = AuthState() def get_auth_state() -> AuthState: """Return the module-level auth state singleton.""" return _auth def reset_auth_state() -> None: """Reset the auth state to its initial (unauthenticated) values.""" global _auth _auth = AuthState() # --------------------------------------------------------------------------- # Callback HTML (adapted from SDK index.html) # --------------------------------------------------------------------------- _CALLBACK_HTML = """\ UiPath CLI Authentication

Authenticate CLI

Completing authentication flow...

✓
×
Processing authentication request...
Authenticating...
Securely exchanging authorization code for access tokens.
""" # --------------------------------------------------------------------------- # Callback HTTP server # --------------------------------------------------------------------------- def _make_handler( html: str, port: int, csrf_state: str, token_callback: Any, ) -> type: """Build the HTTP request handler class with closures over request params.""" class Handler(http.server.BaseHTTPRequestHandler): def log_message(self, fmt: str, *args: Any) -> None: pass def do_GET(self) -> None: content = html.encode() self.send_response(200) self.send_header("Content-Type", "text/html") self.send_header("Content-Length", str(len(content))) self._cors() self.end_headers() self.wfile.write(content) def do_POST(self) -> None: if self.path == f"/set_token/{csrf_state}": length = int(self.headers.get("Content-Length", 0)) try: body = json.loads(self.rfile.read(length)) except (json.JSONDecodeError, ValueError): self.send_error(400, "Malformed JSON") return self.send_response(200) self._cors() self.end_headers() self.wfile.write(b"OK") # Small delay so the browser gets the response before we shut down time.sleep(0.5) token_callback(body) else: self.send_error(404) def do_OPTIONS(self) -> None: self.send_response(200) self._cors() self.end_headers() def _cors(self) -> None: origin = f"http://localhost:{port}" self.send_header("Access-Control-Allow-Origin", origin) self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.send_header("Access-Control-Allow-Headers", "Content-Type") return Handler class _CallbackServer: """Temporary HTTP server that receives the OAuth token from the browser.""" def __init__(self, httpd: socketserver.TCPServer) -> None: self.httpd: socketserver.TCPServer | None = httpd self.port: int = httpd.server_address[1] self._thread: threading.Thread | None = None self._shutdown = False def start(self) -> None: self._shutdown = False self._thread = threading.Thread(target=self._serve, daemon=True) self._thread.start() def _serve(self) -> None: try: while not self._shutdown and self.httpd: self.httpd.handle_request() except Exception: pass def stop(self) -> None: self._shutdown = True if self.httpd: self.httpd.server_close() self.httpd = None # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- _VALID_ENVIRONMENTS = {"cloud", "staging", "alpha"} def build_auth_url(environment: str) -> dict[str, Any]: """Start the OAuth flow: generate PKCE, spin up callback server, return auth URL.""" if environment not in _VALID_ENVIRONMENTS: raise ValueError( f"Invalid environment '{environment}'. Must be one of: {', '.join(sorted(_VALID_ENVIRONMENTS))}" ) auth = get_auth_state() # Stop any previous callback server and cancel any pending wait task if auth._callback_server: auth._callback_server.stop() if auth._wait_task and not auth._wait_task.done(): auth._wait_task.cancel() verifier, challenge = _generate_pkce() state = _generate_state() domain = f"https://{environment}.uipath.com" # Build handler and bind server atomically (no TOCTOU race) # We need the port first to construct the HTML, so we do a two-step: # 1. Build a temporary handler, bind to get the port # 2. Rebuild the handler with the correct HTML, then swap def on_token(token_data: dict[str, Any]) -> None: auth.token_data = token_data if auth._loop and auth._token_event: auth._loop.call_soon_threadsafe(auth._token_event.set) # _make_handler needs the port for CORS, and the HTML needs the port for # redirect_uri. We build a placeholder handler to bind, then rebuild. placeholder = _make_handler("", 0, state, on_token) httpd = _try_bind_server(CANDIDATE_PORTS, placeholder) if httpd is None: raise RuntimeError( f"All callback ports ({', '.join(str(p) for p in CANDIDATE_PORTS)}) are in use" ) port = httpd.server_address[1] redirect_uri = f"http://localhost:{port}/oidc/login" html = ( _CALLBACK_HTML.replace("__STATE__", state) .replace("__CODE_VERIFIER__", verifier) .replace("__REDIRECT_URI__", redirect_uri) .replace("__CLIENT_ID__", CLIENT_ID) .replace("__DOMAIN__", domain) ) # Replace the handler with the real one containing the correct HTML/port httpd.RequestHandlerClass = _make_handler(html, port, state, on_token) query = urlencode( { "client_id": CLIENT_ID, "redirect_uri": redirect_uri, "response_type": "code", "scope": SCOPES, "state": state, "code_challenge": challenge, "code_challenge_method": "S256", } ) auth_url = f"{domain}/identity_/connect/authorize?{query}" # Store state for later (preserve remembered tenant info for re-auth) auth.status = "pending" auth.environment = environment auth._code_verifier = verifier auth._state = state auth._port = port # Set up asyncio event for signalling loop = asyncio.get_running_loop() auth._token_event = asyncio.Event() auth._loop = loop server = _CallbackServer(httpd) server.start() auth._callback_server = server auth._wait_task = asyncio.ensure_future(_wait_for_token(auth)) return {"auth_url": auth_url, "status": "pending"} async def _wait_for_token(auth: AuthState) -> None: """Wait for the callback server to receive the token, then resolve tenants.""" try: if auth._token_event: await asyncio.wait_for(auth._token_event.wait(), timeout=300) except asyncio.TimeoutError: logger.warning("OAuth flow timed out after 5 minutes") auth.status = "unauthenticated" return finally: if auth._callback_server: auth._callback_server.stop() auth._callback_server = None if not auth.token_data: auth.status = "unauthenticated" return # Resolve tenants try: domain = f"https://{auth.environment}.uipath.com" access_token = auth.token_data.get("access_token", "") claims = _parse_jwt_payload(access_token) prt_id = claims.get("prt_id", "") url = f"{domain}/{prt_id}/portal_/api/filtering/leftnav/tenantsAndOrganizationInfo" async with httpx.AsyncClient() as client: resp = await client.get( url, headers={"Authorization": f"Bearer {access_token}"} ) resp.raise_for_status() data = resp.json() auth.tenants = data.get("tenants", []) auth.organization = data.get("organization", {}) tenant_names = [t["name"] for t in auth.tenants] if auth._last_tenant and auth._last_tenant in tenant_names: # Re-auth: auto-select the previously used tenant _finalize_tenant(auth, auth._last_tenant) elif len(auth.tenants) == 1: # Auto-select single tenant _finalize_tenant(auth, auth.tenants[0]["name"]) else: auth.status = "needs_tenant" except Exception: logger.exception("Failed to resolve tenants") auth.status = "unauthenticated" def select_tenant(tenant_name: str) -> dict[str, Any]: """Select a tenant and finalize authentication.""" auth = get_auth_state() tenant = next((t for t in auth.tenants if t["name"] == tenant_name), None) if not tenant: raise ValueError(f"Tenant '{tenant_name}' not found") _finalize_tenant(auth, tenant_name) return {"status": "authenticated", "uipath_url": auth.uipath_url} def _finalize_tenant(auth: AuthState, tenant_name: str) -> None: """Write .env and os.environ with the resolved credentials.""" org_name = auth.organization.get("name", "") domain = f"https://{auth.environment}.uipath.com" uipath_url = f"{domain}/{org_name}/{tenant_name}" access_token = auth.token_data.get("access_token", "") auth.uipath_url = uipath_url auth.status = "authenticated" # Remember for seamless re-auth after expiry auth._last_tenant = tenant_name auth._last_org = dict(auth.organization) auth._last_environment = auth.environment # Update os.environ os.environ["UIPATH_ACCESS_TOKEN"] = access_token os.environ["UIPATH_URL"] = uipath_url # Write/update .env file (preserving comments, blank lines, and ordering) env_path = Path.cwd() / ".env" lines: list[str] = [] updated_keys: set[str] = set() new_values = {"UIPATH_ACCESS_TOKEN": access_token, "UIPATH_URL": uipath_url} if env_path.exists(): with open(env_path) as f: for raw_line in f: stripped = raw_line.strip() if "=" in stripped and not stripped.startswith("#"): key = stripped.split("=", 1)[0] if key in new_values: lines.append(f"{key}={new_values[key]}\n") updated_keys.add(key) continue lines.append(raw_line) # Append any keys that weren't already in the file for key, value in new_values.items(): if key not in updated_keys: lines.append(f"{key}={value}\n") with open(env_path, "w") as f: f.writelines(lines) def logout() -> None: """Clear auth state and env vars.""" auth = get_auth_state() if auth._callback_server: auth._callback_server.stop() os.environ.pop("UIPATH_ACCESS_TOKEN", None) os.environ.pop("UIPATH_URL", None) reset_auth_state() def _check_token_expiry(auth: AuthState) -> None: """Flip status to 'expired' if the current token has expired.""" if auth.status != "authenticated": return access_token = auth.token_data.get("access_token", "") if not access_token: return try: claims = _parse_jwt_payload(access_token) exp = claims.get("exp") if exp is not None and float(exp) < time.time(): auth.status = "expired" except Exception: pass def get_status() -> dict[str, Any]: """Return current auth status for the frontend.""" auth = get_auth_state() _check_token_expiry(auth) result: dict[str, Any] = {"status": auth.status} if auth.status == "needs_tenant": result["tenants"] = [t["name"] for t in auth.tenants] if auth.status in ("authenticated", "expired"): result["uipath_url"] = auth.uipath_url return result def restore_session() -> None: """Check env/.env for existing credentials and restore auth state if valid.""" auth = get_auth_state() if auth.status != "unauthenticated": return # Try os.environ first, then .env file access_token = os.environ.get("UIPATH_ACCESS_TOKEN", "") uipath_url = os.environ.get("UIPATH_URL", "") if not access_token or not uipath_url: # Try reading from .env env_path = Path.cwd() / ".env" if env_path.exists(): env_vars: dict[str, str] = {} with open(env_path) as f: for line in f: line = line.strip() if "=" in line and not line.startswith("#"): key, value = line.split("=", 1) env_vars[key.strip()] = value.strip() access_token = access_token or env_vars.get("UIPATH_ACCESS_TOKEN", "") uipath_url = uipath_url or env_vars.get("UIPATH_URL", "") if not access_token or not uipath_url: return # Check token expiry try: claims = _parse_jwt_payload(access_token) exp = claims.get("exp") if exp is not None and float(exp) < time.time(): logger.debug("Existing token is expired, skipping restore") return except Exception: return # Token is valid — restore state auth.status = "authenticated" auth.uipath_url = uipath_url auth.token_data = {"access_token": access_token} # Ensure os.environ is populated os.environ["UIPATH_ACCESS_TOKEN"] = access_token os.environ["UIPATH_URL"] = uipath_url