icmd-python 1.3.2__py3-none-any.whl → 1.5.8__py3-none-any.whl

icmd/core/client.py

@@ -5,12 +5,16 @@ import webbrowser
 from datetime import UTC, datetime, timedelta
 from getpass import getpass
 from http import HTTPStatus
-from typing import Any
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from pathlib import Path
 
 import requests
 from requests.adapters import HTTPAdapter
 
 from .credentials import CredentialManager
+from .slim_packages import ensure_slim_packages
 
 
 class ICMDAuthenticationError(Exception):
@@ -35,20 +39,50 @@ class ICMD:
     _DEFAULT_TIMEOUT = 30
     _TOKEN_EXPIRE_BUFFER_MS = 5000
     _HTTP_NOT_FOUND = 404
-
-    def __init__(self, domain: str, jupyterlite: bool = False, skip_auth: bool = False):
+    _REFRESH_TIMEOUT = 10
+
+    # Endpoints excluded from auto-refresh retry
+    EXCLUDED_ENDPOINTS: ClassVar[set[str]] = {
+        "account/auth/login",
+        "account/auth/refresh",
+        "account/auth/logout",
+        "account/auth/mfa/verify",
+        "account/auth/saml/login",
+        "account/auth/saml/acs",
+    }
+
+    # Session-wide qtpy_slim tracking (shared across all ICMD instances)
+    # Prevents loading conflicting package versions within same Python session
+    _session_qtpy_version: dict | None = None
+    _session_qtpy_server: str | None = None
+    _session_qtpy_path: "Path | None" = None
+
+    def __init__(
+        self,
+        domain: str,
+        auth_method: str | None = None,
+        developer_token: str | None = None,
+        clear_cache: bool = False,
+        skip_auth: bool = False,
+    ):
         """Initialize ICMD client with required domain parameter.
 
         Args:
             domain: REQUIRED domain (e.g., 'icmd.questek.com', 'cust.icmd.questek.com')
-            jupyterlite: Skip auto-authentication for JupyterLite
+            auth_method: Override auth method detection. Options: "SAML" or "PASSWORD".
+                If not specified, auto-detects based on server config (prefers SAML).
+            developer_token: Developer API token for non-interactive use (scripts, CI/CD).
+                When provided, skips all authentication flows.
+            clear_cache: If True, clears cached credentials before initialization,
+                forcing fresh authentication. Useful for switching accounts or
+                troubleshooting authentication issues.
             skip_auth: Skip auto-authentication (useful for testing or non-interactive
                 environments)
 
         Raises
         ------
         ValueError
-            If domain is not provided or invalid
+            If domain is not provided or invalid, or if auth_method is invalid
         """
         if not domain:
             raise ValueError("Domain is required: ICMD('icmd.questek.com')")
@@ -70,24 +104,57 @@ class ICMD:
         # Set reasonable timeouts by default
         self._default_timeout = self._DEFAULT_TIMEOUT
 
+        # Cache for server config (reduces redundant HTTP requests)
+        self._server_config_cache = None
+
         self._credentials = CredentialManager()
-        self.auth_method = self._auto_detect_auth_method()
-        self.jupyterlite = jupyterlite
+
+        # Clear cached credentials if requested
+        if clear_cache:
+            self._credentials.clear_domain_credentials(self.domain)
+
+        # Handle developer token authentication (non-interactive)
+        if developer_token:
+            self._session.headers["Authorization"] = f"Bearer {developer_token}"
+            self._developer_token = developer_token
+            # Don't fail ICMD initialization if slim package management fails
+            with contextlib.suppress(Exception):
+                ensure_slim_packages(self)
+            return
+
+        self._developer_token = None
+
+        # Validate and set auth method
+        if auth_method:
+            if auth_method.upper() not in ("SAML", "PASSWORD"):
+                raise ValueError(
+                    f"Invalid auth_method: {auth_method}. Must be 'SAML' or 'PASSWORD'"
+                )
+            self.auth_method = auth_method.upper()
+        else:
+            self.auth_method = self._auto_detect_auth_method()
 
         self._refresh_token = ""
-        self._id_token = ""
-        self._id_token_expiration = datetime.now(UTC)
+        self._access_token = ""
+        self._access_token_expiration = datetime.now(UTC)
         self._token_expire_buffer = timedelta(milliseconds=self._TOKEN_EXPIRE_BUFFER_MS)
 
         # Initialize default timeout (needed for tests that bypass __init__)
         if not hasattr(self, "_default_timeout"):
             self._default_timeout = self._DEFAULT_TIMEOUT
 
+        # Display banner message if available (before auth prompts)
+        self._display_banner_message()
+
         self._load_cached_credentials()
 
-        if not (jupyterlite or skip_auth):
+        if not skip_auth:
             self._ensure_authenticated()
 
+            # Don't fail ICMD initialization if slim package management fails
+            with contextlib.suppress(Exception):
+                ensure_slim_packages(self)
+
     def _validate_domain(self, domain: str) -> str:
         """Validate domain format and return normalized domain.
 
@@ -108,14 +175,26 @@ class ICMD:
 
         normalized = self._normalize_domain(domain)
 
+        # Accept localhost for local development
+        # Check both the original domain and the normalized one
+        domain_to_check = domain.replace("http://", "").replace("https://", "")
+        if self._is_local_domain(domain_to_check):
+            return normalized
+
         # Accept: icmd.questek.com or *.icmd.questek.com
         if ".questek.com" in normalized:
             return normalized
 
-        raise ValueError(f"Invalid domain: {domain}. Must be a 'questek.com' subdomain")
+        raise ValueError(
+            f"Invalid domain: {domain}. "
+            "Must be a 'questek.com' subdomain or localhost for development"
+        )
 
     def _normalize_domain(self, domain: str) -> str:
-        """Normalize domain to https:// format."""
+        """Normalize domain with appropriate protocol.
+
+        Uses http:// for localhost, https:// otherwise.
+        """
         if not isinstance(domain, str):
             raise ICMDValidationError("Domain must be a string")
 
@@ -124,25 +203,71 @@ class ICMD:
             raise ICMDValidationError("Domain cannot be empty")
 
         if not domain.startswith("http"):
-            domain = f"https://{domain}"
+            # Use http:// for localhost/local development, https:// for production
+            protocol = "http" if self._is_local_domain(domain) else "https"
+            domain = f"{protocol}://{domain}"
         return domain
 
+    def _is_local_domain(self, domain: str) -> bool:
+        """Check if domain is a local development domain.
+
+        Args
+        ----
+            domain: Domain string to check (without protocol)
+
+        Returns
+        -------
+            True if domain is localhost, 127.0.0.1, or similar local patterns
+        """
+        domain_lower = domain.lower()
+        # Check for localhost and local IP addresses
+        return domain_lower.startswith(("localhost", "127.", "0.0.0.0", "[::1]"))
+
+    def _display_banner_message(self) -> None:
+        """Display banner message from server config if available."""
+        try:
+            config = self._get_server_config()
+            banner_message = config.get("BANNER_MESSAGE", "").strip()
+
+            if banner_message:
+                # Display banner in a noticeable format with colored border
+                border = "=" * 70
+                print(f"\n\033[96m{border}\033[0m")
+                print(f"\033[96m{banner_message}\033[0m")
+                print(f"\033[96m{border}\033[0m\n")
+        except Exception:
+            # Silently fail if banner retrieval fails
+            pass
+
     def _get_server_config(self) -> dict:
-        """Fetch server configuration from /api/v1/public/config/ endpoint."""
+        """Fetch server configuration from /api/v1/public/config/ endpoint.
+
+        Results are cached per instance to avoid redundant HTTP requests during initialization.
+        Empty configs are not cached to allow retry.
+        """
+        # Return cached config if available and non-empty
+        if self._server_config_cache:
+            return self._server_config_cache
+
         try:
             config_url = f"{self.api_root}/public/config/"
             response = self._session.get(config_url, timeout=5)
 
             if response.status_code == HTTPStatus.OK:
-                return response.json()
+                config = response.json()
+                # Only cache non-empty successful responses
+                if config:
+                    self._server_config_cache = config
+                return config
         except (requests.RequestException, ValueError, OSError, Exception):
             # Network issues, JSON parsing errors, timeouts, or any other errors
             pass
 
+        # Don't cache failures or empty responses - allow retry on next call
         return {}
 
     def _enhance_auth_error(self, error_response: requests.Response, auth_method: str) -> str:
-        """Enhance authentication error messages with Cognito-specific context."""
+        """Enhance authentication error messages with Django-specific context."""
         base_error = f"{auth_method} authentication failed"
         enhanced_error = base_error
 
@@ -151,38 +276,38 @@ class ICMD:
             if not isinstance(error_data, dict):
                 return enhanced_error
 
-            # Extract error message
-            error_msg = error_data.get("error") or error_data.get("message", "")
+            # Extract error details from Django response
+            error = error_data.get("error", {})
+            if isinstance(error, dict):
+                error_code = error.get("code", "")
+                error_msg = error.get("message", "")
+            else:
+                error_code = ""
+                error_msg = error_data.get("message", "") or error_data.get("detail", "")
 
-            # Map error types to enhanced messages
+            # Map Django error codes to enhanced messages
             error_enhancements = {
-                "InvalidParameterException": "Invalid credentials or user not found",
-                "UserNotFoundException": "User account not found",
-                "CodeMismatchException": "Invalid authentication code",
-                "ExpiredCodeException": "Authentication code has expired, please try again",
-                "TooManyRequestsException": "Too many attempts, please wait before trying again",
+                "invalid_credentials": "Invalid email or password",
+                "account_inactive": "Account is inactive, contact administrator",
+                "invalid_token": "Invalid or expired token",
+                "invalid_code": "Invalid or expired MFA code",
+                "mfa_required": "Multi-factor authentication required",
+                "mfa_setup_required": (
+                    "MFA not configured. Please configure MFA at "
+                    f"https://{self.domain.replace('https://', '')}/account/security/"
+                ),
+                "rate_limit_exceeded": "Too many attempts, please wait before trying again",
             }
 
-            # Check for specific error types
-            for error_type, enhancement in error_enhancements.items():
-                if error_type in error_msg:
-                    enhanced_error = f"{base_error}: {enhancement}"
-                    break
-            else:
-                # Handle NotAuthorizedException with method-specific messages
-                if "NotAuthorizedException" in error_msg:
-                    if auth_method == "PASSWORD":
-                        enhanced_error = f"{base_error}: Incorrect username or password"
-                    else:
-                        enhanced_error = f"{base_error}: Access denied by identity provider"
-                # Handle networking errors
-                elif "NetworkingError" in error_msg or "timeout" in error_msg.lower():
-                    enhanced_error = (
-                        f"{base_error}: Network connectivity issue with identity provider"
-                    )
-                # Generic error message if available
-                elif error_msg:
-                    enhanced_error = f"{base_error}: {error_msg}"
+            # Check for specific error codes
+            if error_code in error_enhancements:
+                enhanced_error = f"{base_error}: {error_enhancements[error_code]}"
+            # Handle networking errors
+            elif "network" in error_msg.lower() or "timeout" in error_msg.lower():
+                enhanced_error = f"{base_error}: Network connectivity issue"
+            # Generic error message if available
+            elif error_msg:
+                enhanced_error = f"{base_error}: {error_msg}"
 
         except (ValueError, KeyError):
             pass
@@ -194,22 +319,49 @@ class ICMD:
         if cached := self._credentials.get_auth_method(self.domain):
             return cached
 
-        # Server config discovery for Cognito IdP detection
+        # Server config discovery
         config = self._get_server_config()
-        if idps := config.get("identityProviders", []):
-            # Prefer external IdP (Azure AD, etc.) over native Cognito password
-            external_idps = [idp for idp in idps if idp != "UP"]
-            if external_idps:
-                return "SAML"  # External IdP available
-            if "UP" in idps:
-                return "PASSWORD"  # Cognito native user/password
+
+        # Check Django AUTH_METHODS config
+        if auth_methods := config.get("AUTH_METHODS", []):
+            has_saml = any(method.get("type") == "saml" for method in auth_methods)
+            has_password = any(method.get("type") == "password" for method in auth_methods)
+
+            # If both methods available, let user choose
+            if has_saml and has_password:
+                return self._prompt_auth_method_selection()
+
+            if has_saml:
+                return "SAML"
+            if has_password:
+                return "PASSWORD"
 
         # Fallback: When server config is unavailable, default to PASSWORD
-        # NOTE: The consent endpoint exists for both UP-only and SAML-enabled instances,
-        # so we cannot reliably determine SAML availability without server config.
-        # Conservative approach: assume PASSWORD when config is unavailable.
         return "PASSWORD"
 
+    def _prompt_auth_method_selection(self) -> str:
+        """Prompt user to select authentication method when multiple options available."""
+        try:
+            print("\n\033[94mMultiple authentication methods available:\033[0m")
+            print("  [1] SAML SSO  [2] Password")
+
+            while True:
+                try:
+                    choice = input("Select authentication method (1 or 2): ").strip()
+
+                    if choice == "1":
+                        return "SAML"
+                    if choice == "2":
+                        return "PASSWORD"
+                    print("Invalid choice. Please enter 1 or 2.")
+                except KeyboardInterrupt:
+                    # Default to SAML on Ctrl+C
+                    print("\nDefaulting to SAML SSO...")
+                    return "SAML"
+        except EOFError:
+            # Non-interactive environment - default to SAML
+            return "SAML"
+
     @property
     def api_root(self) -> str:
         """Return the root of the ICMD API."""
@@ -228,23 +380,24 @@ class ICMD:
 
     def _is_token_expired(self) -> bool:
         """Check if current token is expired."""
-        return self._id_token_expiration - self._token_expire_buffer <= datetime.now(UTC)
+        return self._access_token_expiration - self._token_expire_buffer <= datetime.now(UTC)
 
     def _ensure_authenticated(self) -> None:
         """Ensure we have a valid authentication token with smart refresh."""
+        # Developer tokens bypass authentication flow
+        if hasattr(self, "_developer_token") and self._developer_token:
+            return
+
         # Check if we already have a valid token
-        if self._id_token and not self._is_token_expired():
+        if self._access_token and not self._is_token_expired():
             return
 
         # Try refresh token if available
         if self._refresh_token:
-            try:
-                self._refresh_access_token()
+            if self._refresh_access_token():
                 return
-            except Exception:
-                # Clear invalid refresh token
-                self._refresh_token = ""
-                self._credentials.clear_refresh_token(self.domain)
+            # Clear invalid refresh token from storage
+            self._credentials.clear_refresh_token(self.domain)
 
         # Full authentication required
         self._perform_authentication()
@@ -253,28 +406,42 @@ class ICMD:
         """Manually trigger authentication. Useful when ICMD was created with skip_auth=True."""
         self._ensure_authenticated()
 
-    def _refresh_access_token(self) -> None:
-        """Refresh access token using refresh token."""
+        # Don't fail if slim package management fails
+        with contextlib.suppress(Exception):
+            ensure_slim_packages(self)
+
+    def _refresh_access_token(self) -> bool:
+        """Refresh access token using refresh token.
+
+        Returns
+        -------
+            bool: True if refresh successful, False otherwise
+        """
         try:
             url = f"{self.api_root}/account/auth/refresh/"
             response = self._session.post(
-                url, json={"refreshToken": self._refresh_token}, timeout=self._default_timeout
+                url,
+                json={"refresh_token": self._refresh_token},
+                timeout=self._REFRESH_TIMEOUT,
             )
 
             if response.status_code != HTTPStatus.OK:
-                error_msg = "Token refresh failed"
-                try:
-                    error_data = response.json()
-                    if isinstance(error_data, dict) and "error" in error_data:
-                        error_msg = f"Token refresh failed: {error_data['error']}"
-                except (ValueError, KeyError):
-                    pass
-                raise ICMDAuthenticationError(error_msg)
+                # Clear tokens on failure
+                self._access_token = ""
+                self._refresh_token = ""
+                self._session.headers.pop("Authorization", None)
+                return False
 
+            # Handle token rotation (new refresh_token in response)
             self._handle_auth_response(response)
+            return True
 
-        except requests.RequestException as e:
-            raise ICMDConnectionError(f"Network error during token refresh: {e}") from e
+        except (requests.RequestException, Exception):
+            # Clear tokens on any error
+            self._access_token = ""
+            self._refresh_token = ""
+            self._session.headers.pop("Authorization", None)
+            return False
 
     def _perform_authentication(self) -> None:
         """Perform full authentication flow."""
@@ -286,22 +453,20 @@ class ICMD:
         self._save_credentials()
 
     def _password_authentication(self) -> None:
-        """Authenticate with username/password."""
+        """Authenticate with email/password (Django JWT)."""
         try:
-            username = input("Your ICMD® email: ").strip()
-            if not username:
-                raise ICMDValidationError("Username cannot be empty")
+            email = input("Your ICMD® email: ").strip()
+            if not email:
+                raise ICMDValidationError("Email cannot be empty")
 
             password = getpass("Your ICMD® password: ")
             if not password:
                 raise ICMDValidationError("Password cannot be empty")
 
-            mfa_code = input("Your MFA Code (if enabled): ").strip()
-
             url = f"{self.api_root}/account/auth/login/"
             response = self._session.post(
                 url,
-                json={"username": username, "password": password},
+                json={"email": email, "password": password},
                 timeout=self._default_timeout,
             )
 
@@ -310,8 +475,10 @@ class ICMD:
                 raise ICMDAuthenticationError(error_msg)
 
             auth_data = response.json()
-            if auth_data.get("challenge"):
-                self._handle_mfa_challenge(auth_data["challenge"], username, mfa_code)
+
+            # Check if MFA is required
+            if auth_data.get("requires_mfa"):
+                self._handle_mfa_challenge(auth_data)
             else:
                 self._handle_auth_response(response)
 
@@ -325,65 +492,176 @@ class ICMD:
         except KeyboardInterrupt:
             raise ICMDAuthenticationError("Authentication cancelled by user") from None
 
-    def _handle_mfa_challenge(self, challenge: dict, username: str, mfa_code: str) -> None:
-        """Handle MFA challenge."""
+    def _handle_mfa_challenge(self, mfa_response: dict) -> None:
+        """Handle Django MFA challenge with retry logic."""
         try:
-            if not mfa_code:
-                mfa_code = input("Your MFA Code: ").strip()
+            # Extract session token and available methods
+            session_token = mfa_response.get("session_token")
+            available_methods = mfa_response.get("available_methods", [])
+
+            if not session_token:
+                raise ICMDAuthenticationError("Invalid MFA response: missing session token")
+
+            # Detect MFA method
+            if "totp" in available_methods:
+                method = "totp"
+                prompt = "Enter MFA code from authenticator app: "
+            elif "email" in available_methods:
+                method = "email"
+                prompt = "MFA code sent to your email. Enter code: "
+            else:
+                method = available_methods[0] if available_methods else "totp"
+                prompt = "Enter MFA code: "
 
-            if not mfa_code:
-                raise ICMDValidationError("MFA code cannot be empty")
+            # Allow up to 3 attempts
+            max_attempts = 3
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    mfa_code = input(prompt).strip()
 
-            url = f"{self.api_root}/account/auth/mfa/"
-            response = self._session.post(
-                url,
-                json={
-                    "mfa_code": mfa_code,
-                    "session": challenge["Session"],
-                    "username": username,
-                    "challenge_name": challenge.get("ChallengeName"),
-                },
-                timeout=self._default_timeout,
-            )
+                    if not mfa_code:
+                        raise ICMDValidationError("MFA code cannot be empty")
 
-            if response.status_code != HTTPStatus.OK:
-                error_msg = "MFA verification failed"
-                try:
+                    url = f"{self.api_root}/account/auth/mfa/verify/"
+
+                    # Set session token in Authorization header (not in body)
+                    headers = {"Authorization": f"Session {session_token}"}
+
+                    response = self._session.post(
+                        url,
+                        json={
+                            "method": method,
+                            "code": mfa_code,
+                        },
+                        headers=headers,
+                        timeout=self._default_timeout,
+                    )
+
+                    if response.status_code == HTTPStatus.OK:
+                        self._handle_auth_response(response)
+                        return
+
+                    # Handle error
                     error_data = response.json()
-                    if isinstance(error_data, dict) and "error" in error_data:
-                        error_msg = f"MFA failed: {error_data['error']}"
-                except (ValueError, KeyError):
-                    pass
-                raise ICMDAuthenticationError(error_msg)
+                    error = error_data.get("error", {})
+                    error_code = error.get("code", "") if isinstance(error, dict) else ""
 
-            self._handle_auth_response(response)
+                    if error_code == "invalid_code" and attempt < max_attempts:
+                        remaining = max_attempts - attempt
+                        print(f"Invalid code. {remaining} attempt(s) remaining.")
+                        continue
+                    error_msg = self._enhance_auth_error(response, "MFA")
+                    raise ICMDAuthenticationError(error_msg)
+
+                except KeyboardInterrupt:
+                    raise ICMDAuthenticationError("MFA verification cancelled by user") from None
+
+            raise ICMDAuthenticationError("MFA verification failed: Maximum attempts exceeded")
 
         except requests.RequestException as e:
             raise ICMDConnectionError(f"Network error during MFA verification: {e}") from e
-        except KeyboardInterrupt:
-            raise ICMDAuthenticationError("MFA verification cancelled by user") from None
+
+    def _get_saml_provider(self) -> dict:
+        """Get SAML provider configuration from server.
+
+        Returns
+        -------
+            dict: SAML provider info with 'id', 'name', and 'login_url'
+
+        Raises
+        ------
+            ICMDAuthenticationError: If no SAML providers configured or selection fails
+        """
+        try:
+            config = self._get_server_config()
+            saml_providers = config.get("SAML_PROVIDERS", [])
+
+            if not saml_providers:
+                raise ICMDAuthenticationError(
+                    "SAML authentication is not configured on this server. "
+                    "Please contact your administrator or use password authentication."
+                )
+
+            # If only one provider, auto-select
+            if len(saml_providers) == 1:
+                provider = saml_providers[0]
+                print(f"Using SAML provider: {provider.get('name', provider['id'])}")
+                return provider
+
+            # Multiple providers - let user choose
+            print("\n\033[94mMultiple SAML providers available:\033[0m")
+            for idx, provider in enumerate(saml_providers, 1):
+                provider_name = provider.get("name", provider["id"])
+                description = provider.get("description", "")
+                if description:
+                    print(f"  [{idx}] {provider_name} - {description}")
+                else:
+                    print(f"  [{idx}] {provider_name}")
+
+            n_providers = len(saml_providers)
+            while True:
+                try:
+                    choice = input(f"\nSelect SAML provider (1-{n_providers}): ").strip()
+
+                    try:
+                        choice_idx = int(choice) - 1
+                        if 0 <= choice_idx < n_providers:
+                            return saml_providers[choice_idx]
+                        print(
+                            f"Invalid choice. Please enter a number between 1 and {n_providers}."
+                        )
+                    except ValueError:
+                        print(f"Invalid input. Please enter a number between 1 and {n_providers}.")
+
+                except KeyboardInterrupt:
+                    # Default to first provider on Ctrl+C
+                    print(
+                        f"\nDefaulting to {
+                            saml_providers[0].get('name', saml_providers[0]['id'])
+                        }..."
+                    )
+                    return saml_providers[0]
+
+        except EOFError:
+            # Non-interactive environment - use first provider
+            if saml_providers:
+                return saml_providers[0]
+            raise ICMDAuthenticationError(
+                "Cannot select SAML provider in non-interactive environment"
+            ) from None
 
     def _saml_authentication(self) -> None:
-        """Authenticate with SAML SSO."""
+        """Authenticate with Django SAML SSO."""
         try:
-            print("\n\033[94mSigning into ICMD® with SAML SSO...\033[0m\n")
+            from urllib.parse import quote
+
+            # Get SAML provider configuration
+            provider = self._get_saml_provider()
+
+            # Use provider's login URL with SDK flag in relay_state
+            relay_state = quote("/auth/callback?sdk=true")
+            consent_url = f"{self.domain}{provider['login_url']}&relay_state={relay_state}"
 
-            consent_url = f"{self.api_root}/account/auth/consent"
-            print(f"Opening browser: {consent_url}")
-            print("If it doesn't open, copy this URL to your browser:")
-            print(f"\n{consent_url}\n")
+            # Display the authentication link
+            print(f"\n\033[94mSAML SSO Authentication Link:\033[0m\n{consent_url}")
 
+            # Compact messaging for Jupyter notebooks
+            print("\n\033[94mOpening browser...\033[0m Copy the auth code (5 min expiry)")
+
+            # Attempt to open browser automatically
             with contextlib.suppress(Exception):
                 webbrowser.open(consent_url)
 
-            auth_code = input("Please enter the authentication code: ").strip()
+            # Compact input prompt
+            auth_code = input("Auth Code: ").strip()
             if not auth_code:
                 raise ICMDValidationError("Authentication code required")
 
+            # Exchange auth code for tokens
             url = f"{self.api_root}/account/auth/login/"
             response = self._session.post(
                 url,
-                json={"code": auth_code, "redirect_uri": f"{self.api_root}/account/auth/exchange"},
+                json={"auth_code": auth_code},
                 timeout=self._default_timeout,
             )
 
@@ -392,7 +670,7 @@ class ICMD:
                 raise ICMDAuthenticationError(error_msg)
 
             self._handle_auth_response(response)
-            print("\n\033[92mSigned in successfully.\033[0m\n")
+            print("\033[92m✓ Authenticated\033[0m")
 
         except requests.RequestException as e:
             raise ICMDConnectionError(f"Network error during SAML authentication: {e}") from e
@@ -400,28 +678,31 @@ class ICMD:
             raise ICMDAuthenticationError("SAML authentication cancelled by user") from None
 
     def _handle_auth_response(self, response: requests.Response) -> None:
-        """Handle authentication response and extract tokens."""
+        """Handle Django JWT authentication response and extract tokens."""
         auth_data = response.json()
 
-        # Update session headers
-        if user_context := auth_data.get("userFingerprint"):
-            self._session.headers["X-User-Context"] = user_context
+        # Check for error response
+        if not auth_data.get("success", False):
+            error = auth_data.get("error", {})
+            error_msg = error.get("message", "Authentication failed")
+            raise ICMDAuthenticationError(error_msg)
+
+        # Update session headers with user UUID
+        if (user_data := auth_data.get("user")) and (user_uuid := user_data.get("uuid")):
+            self._session.headers["X-User-Context"] = user_uuid
 
         # Extract tokens
-        self._id_token = auth_data.get("idToken", "")
-        if refresh_token := auth_data.get("refreshToken"):
+        self._access_token = auth_data.get("access_token", "")
+        if refresh_token := auth_data.get("refresh_token"):
             self._refresh_token = refresh_token
 
-        # Set expiration
-        if expires_at := auth_data.get("expiresAt"):
-            self._id_token_expiration = datetime.fromisoformat(expires_at)
-            # Ensure timezone awareness for comparison
-            if self._id_token_expiration.tzinfo is None:
-                self._id_token_expiration = self._id_token_expiration.replace(tzinfo=UTC)
+        # Calculate expiration from expires_in (seconds)
+        if expires_in := auth_data.get("expires_in"):
+            self._access_token_expiration = datetime.now(UTC) + timedelta(seconds=expires_in)
 
         # Set bearer auth
-        if self._id_token:
-            self._session.headers["Authorization"] = f"Bearer {self._id_token}"
+        if self._access_token:
+            self._session.headers["Authorization"] = f"Bearer {self._access_token}"
 
         # Handle cookies
         for cookie in response.cookies:
@@ -448,15 +729,35 @@ class ICMD:
             response = self._session.request(method, url, **kwargs)
 
             # Handle token expiration with single retry
-            if response.status_code == HTTPStatus.UNAUTHORIZED and not kwargs.get("_retry_auth"):
-                # Clear token and re-authenticate
-                self._id_token = ""
-                self._session.headers.pop("Authorization", None)
-                self._ensure_authenticated()
-
-                # Retry once with fresh token
-                kwargs["_retry_auth"] = True
-                response = self._session.request(method, url, **kwargs)
+            if response.status_code == HTTPStatus.UNAUTHORIZED and not kwargs.get(
+                "_retry_attempted"
+            ):
+                # Check if endpoint is excluded from auto-refresh
+                if endpoint in self.EXCLUDED_ENDPOINTS:
+                    return response
+
+                # Check error detail for selective logout
+                try:
+                    error_data = response.json()
+                    error_detail = error_data.get("detail", "").lower()
+                    # Only auto-refresh for specific token-related errors
+                    if any(
+                        msg in error_detail
+                        for msg in ["invalid token", "token has expired", "not authenticated"]
+                    ):
+                        # Set retry flag to prevent infinite loop
+                        kwargs["_retry_attempted"] = True
+
+                        # Attempt token refresh
+                        if self._refresh_access_token():
+                            # Retry with refreshed token
+                            response = self._session.request(method, url, **kwargs)
+                        else:
+                            # Refresh failed, return original 401 response
+                            return response
+                except (ValueError, KeyError):
+                    # If response is not JSON or doesn't have expected structure, don't retry
+                    pass
 
             return response
 
@@ -483,3 +784,32 @@ class ICMD:
     def delete(self, endpoint: str, **kwargs) -> requests.Response:
         """DELETE request."""
         return self.request("DELETE", endpoint, **kwargs)
+
+    def logout(self, clear_saved_credentials: bool = False) -> None:
+        """Logout and clear session.
+
+        Args:
+            clear_saved_credentials: If True, also clear cached credentials from disk.
+                                   Default False (preserves refresh token for faster re-auth).
+        """
+        # Best effort logout call to server (ignore errors)
+        try:
+            if self._access_token:
+                url = f"{self.api_root}/account/auth/logout/"
+                self._session.post(url, timeout=5)
+        except Exception:
+            pass  # Ignore errors, proceed with local cleanup
+
+        # Clear in-memory tokens
+        self._access_token = ""
+        self._refresh_token = ""
+        self._access_token_expiration = datetime.now(UTC)
+
+        # Remove Authorization header
+        self._session.headers.pop("Authorization", None)
+
+        # Optionally clear saved credentials
+        if clear_saved_credentials:
+            self._credentials.clear_refresh_token(self.domain)
+
+        print("Logged out successfully.")