feat(auth): add support for custom urls

Author: ionmincu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.1.23"
+version = "2.1.24"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"

src/uipath/_cli/__init__.py

@@ -3,7 +3,7 @@
 
 import click
 
-from .cli_auth import auth as auth  # type: ignore
+from .cli_auth import auth as auth
 from .cli_deploy import deploy as deploy  # type: ignore
 from .cli_eval import eval as eval  # type: ignore
 from .cli_init import init as init  # type: ignore

src/uipath/_cli/_auth/_oidc_utils.py

@@ -65,5 +65,8 @@ def get_auth_url(domain: str) -> tuple[str, str, str]:
     }
 
     query_string = urlencode(query_params)
-    url = f"https://{domain}.uipath.com/identity_/connect/authorize?{query_string}"
+    if domain.startswith("http"):
+        url = f"{domain}/identity_/connect/authorize?{query_string}"
+    else:
+        url = f"https://{domain}.uipath.com/identity_/connect/authorize?{query_string}"
     return url, code_verifier, state

src/uipath/_cli/_auth/_portal_service.py

@@ -65,7 +65,10 @@ def get_tenants_and_organizations(self) -> TenantsAndOrganizationInfoResponse:
         if self._client is None:
             raise RuntimeError("HTTP client is not initialized")
 
-        url = f"https://{self.domain}.uipath.com/{self.prt_id}/portal_/api/filtering/leftnav/tenantsAndOrganizationInfo"
+        if self.domain and self.domain.startswith("http"):
+            url = f"{self.domain}/{self.prt_id}/portal_/api/filtering/leftnav/tenantsAndOrganizationInfo"
+        else:
+            url = f"https://{self.domain}.uipath.com/{self.prt_id}/portal_/api/filtering/leftnav/tenantsAndOrganizationInfo"
         response = self._client.get(
             url, headers={"Authorization": f"Bearer {self.access_token}"}
         )
@@ -213,12 +216,19 @@ def select_tenant(
     account_name = tenants_and_organizations["organization"]["name"]
     console.info(f"Selected tenant: {click.style(tenant_name, fg='cyan')}")
 
+    if domain.startswith("http"):
+        base_url = domain
+    else:
+        base_url = f"https://{domain if domain else 'cloud'}.uipath.com"
+
+    uipath_url = f"{base_url}/{account_name}/{tenant_name}"
+
     update_env_file(
         {
-            "UIPATH_URL": f"https://{domain if domain else 'alpha'}.uipath.com/{account_name}/{tenant_name}",
+            "UIPATH_URL": uipath_url,
             "UIPATH_TENANT_ID": tenants_and_organizations["tenants"][tenant_idx]["id"],
             "UIPATH_ORGANIZATION_ID": tenants_and_organizations["organization"]["id"],
         }
     )
 
-    return f"https://{domain if domain else 'alpha'}.uipath.com/{account_name}/{tenant_name}"
+    return uipath_url

src/uipath/_cli/_auth/index.html

@@ -477,7 +477,12 @@ <h1 class="auth-title" id="main-title">Authenticate CLI</h1>
                 formData.append('client_id', '__PY_REPLACE_CLIENT_ID__');
                 formData.append('code_verifier', codeVerifier);
 
-                const response = await fetch('https://__PY_REPLACE_DOMAIN__.uipath.com/identity_/connect/token', {
+                const domain = '__PY_REPLACE_DOMAIN__';
+                const tokenUrl = domain.startsWith('http')
+                    ? `${domain}/identity_/connect/token`
+                    : `https://${domain}.uipath.com/identity_/connect/token`;
+
+                const response = await fetch(tokenUrl, {
                     method: 'POST',
                     headers: {
                         'Content-Type': 'application/x-www-form-urlencoded'

src/uipath/_cli/cli_auth.py

@@ -1,9 +1,9 @@
-# type: ignore
 import asyncio
 import json
 import os
 import socket
 import webbrowser
+from urllib.parse import urlparse
 
 import click
 from dotenv import load_dotenv
@@ -33,10 +33,10 @@ def is_port_in_use(port: int) -> bool:
 
 def set_port():
     auth_config = get_auth_config()
-    port = auth_config.get("port", 8104)
-    port_option_one = auth_config.get("portOptionOne", 8104)
-    port_option_two = auth_config.get("portOptionTwo", 8055)
-    port_option_three = auth_config.get("portOptionThree", 42042)
+    port = int(auth_config.get("port", 8104))
+    port_option_one = int(auth_config.get("portOptionOne", 8104))  # type: ignore
+    port_option_two = int(auth_config.get("portOptionTwo", 8055))  # type: ignore
+    port_option_three = int(auth_config.get("portOptionThree", 42042))  # type: ignore
     if is_port_in_use(port):
         if is_port_in_use(port_option_one):
             if is_port_in_use(port_option_two):
@@ -85,12 +85,15 @@ def set_port():
 def auth(
     domain,
     force: None | bool = False,
-    client_id: str = None,
-    client_secret: str = None,
-    base_url: str = None,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    base_url: str | None = None,
 ):
     """Authenticate with UiPath Cloud Platform.
 
+    The domain for authentication is determined by the UIPATH_URL environment variable if set.
+    Otherwise, it can be specified with --cloud (default), --staging, or --alpha flags.
+
     Interactive mode (default): Opens browser for OAuth authentication.
     Unattended mode: Use --client-id, --client-secret and --base-url for client credentials flow.
 
@@ -99,6 +102,17 @@ def auth(
     - Set REQUESTS_CA_BUNDLE to specify a custom CA bundle for SSL verification
     - Set UIPATH_DISABLE_SSL_VERIFY to disable SSL verification (not recommended)
     """
+    uipath_url = os.getenv("UIPATH_URL")
+    if uipath_url and domain == "cloud":  # "cloud" is the default
+        parsed_url = urlparse(uipath_url)
+        if parsed_url.scheme and parsed_url.netloc:
+            domain = f"{parsed_url.scheme}://{parsed_url.netloc}"
+        else:
+            console.error(
+                f"Malformed UIPATH_URL: '{uipath_url}'. Please ensure it includes both scheme and netloc (e.g., 'https://cloud.uipath.com')."
+            )
+            return
+
     # Check if client credentials are provided for unattended authentication
     if client_id and client_secret:
         if not base_url:

tests/cli/test_auth.py

@@ -0,0 +1,141 @@
+import os
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from click.testing import CliRunner
+
+from uipath._cli.cli_auth import auth
+
+"""
+Unit tests for the 'uipath auth' command.
+
+This test suite covers the following scenarios for the authentication logic:
+
+1.  **UIPATH_URL Environment Variable**: 
+    Ensures the `auth` command correctly uses the domain from the `UIPATH_URL` 
+    environment variable when no specific environment flag is used.
+
+2.  **--alpha Flag**: 
+    Verifies that the `--alpha` flag correctly uses the 'alpha' environment and 
+    overrides the `UIPATH_URL` environment variable if it is set.
+
+3.  **--staging Flag**: 
+    Verifies that the `--staging` flag correctly uses the 'staging' environment and 
+    overrides the `UIPATH_URL` environment variable if it is set.
+
+4.  **--cloud Flag**: 
+    Checks that the `--cloud` flag works as expected, using the 'cloud' environment,
+    when no `UIPATH_URL` environment variable is set.
+
+5.  **Default Behavior**: 
+    Confirms that the command defaults to the 'cloud' environment when no flags 
+    or environment variables are provided.
+"""
+
+
+@pytest.mark.parametrize(
+    "scenario_name, cli_args, env_vars, expected_url_part, expected_select_tenant_return",
+    [
+        (
+            "auth_with_uipath_url_env_variable",
+            ["--force"],
+            {"UIPATH_URL": "https://custom.automationsuite.org/org/tenant"},
+            "https://custom.automationsuite.org/identity_/connect/authorize",
+            "https://custom.automationsuite.org/DefaultOrg/DefaultTenant",
+        ),
+        (
+            "auth_with_uipath_url_env_variable_with_trailing_slash",
+            ["--force"],
+            {"UIPATH_URL": "https://custom.uipath.com/org/tenant/"},
+            "https://custom.uipath.com/identity_/connect/authorize",
+            "https://custom.uipath.com/DefaultOrg/DefaultTenant",
+        ),
+        (
+            "auth_with_alpha_flag",
+            ["--alpha", "--force"],
+            {"UIPATH_URL": "https://custom.uipath.com/org/tenant"},
+            "https://alpha.uipath.com/identity_/connect/authorize",
+            "https://alpha.uipath.com/DefaultOrg/DefaultTenant",
+        ),
+        (
+            "auth_with_staging_flag",
+            ["--staging", "--force"],
+            {"UIPATH_URL": "https://custom.uipath.com/org/tenant"},
+            "https://staging.uipath.com/identity_/connect/authorize",
+            "https://staging.uipath.com/DefaultOrg/DefaultTenant",
+        ),
+        (
+            "auth_with_cloud_flag",
+            ["--cloud", "--force"],
+            {},
+            "https://cloud.uipath.com/identity_/connect/authorize",
+            "https://cloud.uipath.com/DefaultOrg/DefaultTenant",
+        ),
+        (
+            "auth_default_to_cloud",
+            ["--force"],
+            {},
+            "https://cloud.uipath.com/identity_/connect/authorize",
+            "https://cloud.uipath.com/DefaultOrg/DefaultTenant",
+        ),
+    ],
+    ids=[
+        "uipath_url_env",
+        "alpha_flag_overrides_env",
+        "staging_flag_overrides_env",
+        "cloud_flag",
+        "default_to_cloud",
+        "uipath_url_env_with_trailing_slash",
+    ],
+)
+def test_auth_scenarios(
+    scenario_name, cli_args, env_vars, expected_url_part, expected_select_tenant_return
+):
+    """
+    Test 'uipath auth' with different configurations.
+    """
+    runner = CliRunner()
+    with (
+        patch("uipath._cli.cli_auth.webbrowser.open") as mock_open,
+        patch("uipath._cli.cli_auth.HTTPServer") as mock_server,
+        patch("uipath._cli.cli_auth.PortalService") as mock_portal_service,
+    ):
+        mock_server.return_value.start = AsyncMock(
+            return_value={"access_token": "test_token"}
+        )
+        mock_portal_service.return_value.__enter__.return_value.get_tenants_and_organizations.return_value = {
+            "tenants": [{"name": "DefaultTenant", "id": "tenant-id"}],
+            "organization": {"name": "DefaultOrg", "id": "org-id"},
+        }
+        mock_portal_service.return_value.__enter__.return_value.select_tenant.return_value = expected_select_tenant_return
+
+        with runner.isolated_filesystem():
+            for key, value in env_vars.items():
+                os.environ[key] = value
+
+            result = runner.invoke(auth, cli_args)
+
+            for key in env_vars:
+                del os.environ[key]
+
+            assert result.exit_code == 0, (
+                f"Scenario '{scenario_name}' failed with exit code {result.exit_code}: {result.output}"
+            )
+            mock_open.assert_called_once()
+            call_args = mock_open.call_args[0][0]
+            assert expected_url_part in call_args
+
+
+def test_auth_with_malformed_url():
+    """
+    Test that 'uipath auth' handles a malformed UIPATH_URL gracefully.
+    """
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        os.environ["UIPATH_URL"] = "custom.uipath.com"
+        result = runner.invoke(auth, ["--force"])
+        del os.environ["UIPATH_URL"]
+
+        assert result.exit_code == 1
+        assert "Malformed UIPATH_URL" in result.output
+        assert "custom.uipath.com" in result.output