Added support for security token authentication (#231)

Author: lu-ohai

ads/common/auth.py

@@ -5,8 +5,10 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import copy
+from datetime import datetime
 import os
 from dataclasses import dataclass
+import time
 from typing import Any, Callable, Dict, Optional, Union
 
 import ads.telemetry
@@ -17,11 +19,18 @@
 from oci.config import DEFAULT_LOCATION  # "~/.oci/config"
 from oci.config import DEFAULT_PROFILE  # "DEFAULT"
 
+SECURITY_TOKEN_LEFT_TIME = 600
+
+
+class SecurityTokenError(Exception):  # pragma: no cover
+    pass
+
 
 class AuthType(str, metaclass=ExtendedEnumMeta):
     API_KEY = "api_key"
     RESOURCE_PRINCIPAL = "resource_principal"
     INSTANCE_PRINCIPAL = "instance_principal"
+    SECURITY_TOKEN = "security_token"
 
 
 class SingletonMeta(type):
@@ -140,6 +149,15 @@ def set_auth(
 
     >>> ads.set_auth("instance_principal")  # Set instance principal authentication
 
+    >>> ads.set_auth("security_token")  # Set security token authentication
+
+    >>> config = dict(
+    ...     region=us-ashburn-1,
+    ...     key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem,
+    ...     security_token_file=~/.oci/sessions/DEFAULT/token
+    ... )
+    >>> ads.set_auth("security_token", config=config) # Set security token authentication from provided config
+
     >>> singer = oci.signer.Signer(
     ...     user=ocid1.user.oc1..<unique_ID>,
     ...     fingerprint=<fingerprint>,
@@ -274,6 +292,50 @@ def resource_principal(
     return signer_generator(signer_args).create_signer()
 
 
+def security_token(
+    oci_config: Union[str, Dict] = os.path.expanduser(DEFAULT_LOCATION),
+    profile: str = DEFAULT_PROFILE,
+    client_kwargs: Dict = None,
+) -> Dict:
+    """
+    Prepares authentication and extra arguments necessary for creating clients for different OCI services using Security Token.
+
+    Parameters
+    ----------
+    oci_config: Optional[Union[str, Dict]], default is ~/.oci/config
+        OCI authentication config file location or a dictionary with config attributes.
+    profile: Optional[str], is DEFAULT_PROFILE, which is 'DEFAULT'
+        Profile name to select from the config file.
+    client_kwargs: Optional[Dict], default None
+        kwargs that are required to instantiate the Client if we need to override the defaults.
+
+    Returns
+    -------
+    dict
+        Contains keys - config, signer and client_kwargs.
+
+        - The config contains the config loaded from the configuration loaded from `oci_config`.
+        - The signer contains the signer object created from the security token.
+        - client_kwargs contains the `client_kwargs` that was passed in as input parameter.
+
+    Examples
+    --------
+    >>> from ads.common import oci_client as oc
+    >>> auth = ads.auth.security_token(oci_config="/home/datascience/.oci/config", profile="TEST", client_kwargs={"timeout": 6000})
+    >>> oc.OCIClientFactory(**auth).object_storage # Creates Object storage client with timeout set to 6000 using Security Token authentication
+    """
+    signer_args = dict(
+        oci_config=oci_config if isinstance(oci_config, Dict) else {},
+        oci_config_location=oci_config
+        if isinstance(oci_config, str)
+        else os.path.expanduser(DEFAULT_LOCATION),
+        oci_key_profile=profile,
+        client_kwargs=client_kwargs,
+    )
+    signer_generator = AuthFactory().signerGenerator(AuthType.SECURITY_TOKEN)
+    return signer_generator(signer_args).create_signer()
+
+
 def create_signer(
     auth_type: Optional[str] = AuthType.API_KEY,
     oci_config_location: Optional[str] = DEFAULT_LOCATION,
@@ -346,6 +408,12 @@ def create_signer(
     >>> signer_callable = oci.auth.signers.InstancePrincipalsSecurityTokenSigner
     >>> signer_kwargs = dict(log_requests=True) # will log the request url and response data when retrieving
     >>> auth = ads.auth.create_signer(signer_callable=signer_callable, signer_kwargs=signer_kwargs) # instance principals authentication dictionary created based on callable with kwargs parameters
+    >>> config = dict(
+    ...     region=us-ashburn-1,
+    ...     key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem,
+    ...     security_token_file=~/.oci/sessions/DEFAULT/token
+    ... )
+    >>> auth = ads.auth.create_signer(auth_type="security_token", config=config) # security token authentication created based on provided config
     """
     if signer or signer_callable:
         configuration = ads.telemetry.update_oci_client_config()
@@ -365,8 +433,6 @@ def create_signer(
             oci_config=config,
             client_kwargs=client_kwargs,
         )
-        if config:
-            auth_type = AuthType.API_KEY
 
         signer_generator = AuthFactory().signerGenerator(auth_type)
 
@@ -678,6 +744,162 @@ def create_signer(self) -> Dict:
         return signer_dict
 
 
+class SecurityToken(AuthSignerGenerator):
+    """
+    Creates security token auth instance. This signer is intended to be used when signing requests for
+    a given user - it requires that user's private key and security token.
+    It prepares extra arguments necessary for creating clients for variety of OCI services.
+    """
+    SECURITY_TOKEN_GENERIC_HEADERS = [
+        "date",
+        "(request-target)",
+        "host"
+    ]
+    SECURITY_TOKEN_BODY_HEADERS = [
+        "content-length",
+        "content-type",
+        "x-content-sha256"
+    ]
+    SECURITY_TOKEN_REQUIRED = [
+        "security_token_file",
+        "key_file",
+        "region"
+    ]
+
+    def __init__(self, args: Optional[Dict] = None):
+        """
+        Signer created based on args provided. If not provided current values of according arguments
+        will be used from current global state from AuthState class.
+
+        Parameters
+        ----------
+        args: dict
+            args that are required to create Security Token signer. Contains keys: oci_config,
+            oci_config_location, oci_key_profile, client_kwargs.
+
+            - oci_config is a configuration dict that can be used to create clients
+            - oci_config_location - path to config file
+            - oci_key_profile - the profile to load from config file
+            - client_kwargs - optional parameters for OCI client creation in next steps
+        """
+        self.oci_config = args.get("oci_config")
+        self.oci_config_location = args.get("oci_config_location")
+        self.oci_key_profile = args.get("oci_key_profile")
+        self.client_kwargs = args.get("client_kwargs")
+
+    def create_signer(self) -> Dict:
+        """
+        Creates security token configuration and signer with extra arguments necessary for creating clients.
+        Signer constructed from the `oci_config` provided. If not 'oci_config', configuration will be
+        constructed from 'oci_config_location' and 'oci_key_profile' in place.
+
+        Returns
+        -------
+        dict
+            Contains keys - config, signer and client_kwargs.
+
+            - config contains the configuration information
+            - signer contains the signer object created. It is instantiated from signer_callable, or
+            signer provided in args used, or instantiated in place
+            - client_kwargs contains the `client_kwargs` that was passed in as input parameter
+
+        Examples
+        --------
+        >>> signer_args = dict(
+        ...     client_kwargs=client_kwargs
+        ... )
+        >>> signer_generator = AuthFactory().signerGenerator(AuthType.SECURITY_TOKEN)
+        >>> signer_generator(signer_args).create_signer()
+        """
+        if self.oci_config:
+            configuration = ads.telemetry.update_oci_client_config(self.oci_config)
+        else:
+            configuration = ads.telemetry.update_oci_client_config(
+                oci.config.from_file(self.oci_config_location, self.oci_key_profile)
+            )
+
+        logger.info(f"Using 'security_token' authentication.")
+
+        for parameter in self.SECURITY_TOKEN_REQUIRED:
+            if parameter not in configuration:
+                raise ValueError(
+                    f"Parameter `{parameter}` must be provided for using `security_token` authentication."
+                )
+
+        self._validate_and_refresh_token(configuration)
+
+        return {
+            "config": configuration,
+            "signer": oci.auth.signers.SecurityTokenSigner(
+                token=self._read_security_token_file(configuration.get("security_token_file")),
+                private_key=oci.signer.load_private_key_from_file(
+                    configuration.get("key_file"), configuration.get("pass_phrase")
+                ),
+                generic_headers=configuration.get("generic_headers", self.SECURITY_TOKEN_GENERIC_HEADERS),
+                body_headers=configuration.get("body_headers", self.SECURITY_TOKEN_BODY_HEADERS)
+            ),
+            "client_kwargs": self.client_kwargs,
+        }
+
+    def _validate_and_refresh_token(self, configuration: Dict[str, Any]):
+        """Validates and refreshes security token.
+
+        Parameters
+        ----------
+        configuration: Dict
+            Security token configuration.
+        """
+        security_token = self._read_security_token_file(configuration.get("security_token_file"))
+        security_token_container = oci.auth.security_token_container.SecurityTokenContainer(
+            session_key_supplier=None,
+            security_token=security_token
+        )
+
+        if not security_token_container.valid():
+            raise SecurityTokenError(
+                "Security token has expired. Call `oci session authenticate` to generate new session."
+            )
+        
+        time_now = int(time.time())
+        time_expired = security_token_container.get_jwt()["exp"]
+        if time_expired - time_now < SECURITY_TOKEN_LEFT_TIME:
+            if not self.oci_config_location:
+                logger.warning("Can not auto-refresh token. Specify parameter `oci_config_location` through ads.set_auth() or ads.auth.create_signer().")
+            else:
+                result = os.system(f"oci session refresh --config-file {self.oci_config_location} --profile {self.oci_key_profile}")
+                if result == 1:
+                    logger.warning(
+                        "Some error happened during auto-refreshing the token. Continue using the current one that's expiring in less than {SECURITY_TOKEN_LEFT_TIME} seconds."
+                        "Please follow steps in https://docs.oracle.com/en-us/iaas/Content/API/SDKDocs/clitoken.htm to renew token."
+                    )
+        
+        date_time = datetime.fromtimestamp(time_expired).strftime("%Y-%m-%d %H:%M:%S")
+        logger.info(f"Session is valid until {date_time}.")
+
+    def _read_security_token_file(self, security_token_file: str) -> str:
+        """Reads security token from file.
+
+        Parameters
+        ----------
+        security_token_file: str
+            The path to security token file.
+
+        Returns
+        -------
+        str:
+            Security token string.
+        """
+        if not os.path.isfile(security_token_file):
+            raise ValueError("Invalid `security_token_file`. Specify a valid path.")
+        try:
+            token = None
+            with open(security_token_file, 'r') as f:
+                token = f.read()
+            return token
+        except:
+            raise
+
+
 class AuthFactory:
     """
     AuthFactory class which contains list of registered signers and alllows to register new signers.
@@ -687,12 +909,14 @@ class AuthFactory:
         * APIKey
         * ResourcePrincipal
         * InstancePrincipal
+        * SecurityToken
     """
 
     classes = {
         AuthType.API_KEY: APIKey,
         AuthType.RESOURCE_PRINCIPAL: ResourcePrincipal,
         AuthType.INSTANCE_PRINCIPAL: InstancePrincipal,
+        AuthType.SECURITY_TOKEN: SecurityToken,
     }
 
     @classmethod
@@ -726,7 +950,7 @@ def signerGenerator(self, iam_type: Optional[str] = "api_key"):
 
         Returns
         -------
-        :class:`APIKey` or :class:`ResourcePrincipal` or :class:`InstancePrincipal`
+        :class:`APIKey` or :class:`ResourcePrincipal` or :class:`InstancePrincipal` or :class:`SecurityToken`
             returns one of classes, which implements creation of signer of specified type
 
         Raises

ads/opctl/cli.py

@@ -230,7 +230,7 @@ def init_vscode(**kwargs):
         "--auth",
         "-a",
         help="authentication method",
-        type=click.Choice(["api_key", "resource_principal"]),
+        type=click.Choice(["api_key", "resource_principal", "security_token"]),
         default=None,
     ),
 ]

ads/opctl/config/merger.py

@@ -115,7 +115,7 @@ def _fill_config_with_defaults(self, ads_config_path: str) -> None:
             else:
                 self.config["execution"]["auth"] = AuthType.API_KEY
         # determine profile
-        if self.config["execution"]["auth"] != AuthType.API_KEY:
+        if self.config["execution"]["auth"] == AuthType.RESOURCE_PRINCIPAL:
             profile = self.config["execution"]["auth"].upper()
             exec_config.pop("oci_profile", None)
             self.config["execution"]["oci_profile"] = None

docs/source/user_guide/cli/authentication.rst

@@ -62,13 +62,28 @@ You can choose to use the instance principal to authenticate while using the Acc
   mc = ModelCatalog(compartment_id="<compartment_id>")
   mc.list_models()
 
+4. Authenticating Using Security Token
+--------------------------------------
 
-4. Overriding Defaults
+**Prerequisite**
+
+* You have setup security token as per the instruction `here <https://docs.oracle.com/en-us/iaas/Content/API/SDKDocs/clitoken.htm>`_
+
+You can choose to use the security token to authenticate while using the Accelerated Data Science (ADS) SDK by running ``ads.set_auth(auth='security_token')``. For example:
+
+.. code-block:: python
+
+  import ads
+  ads.set_auth(auth='security_token')
+  mc = ModelCatalog(compartment_id="<compartment_id>")
+  mc.list_models()
+
+5. Overriding Defaults
 ----------------------
 
 The default authentication that is used by ADS is set with the ``set_auth()`` method. However, each relevant ADS method has an optional parameter to specify the authentication method to use. The most common use case for this is when you have different permissions in different API keys or there are differences between the permissions granted in the resource principals and your API keys.
 
-By default, ADS uses API keys to sign requests to OCI resources. The ``set_auth()`` method is used to explicitly set a default signing method. This method accepts one of three strings ``"api_key"``, ``"resource_principal"``, or ``instance_principal``.
+By default, ADS uses API keys to sign requests to OCI resources. The ``set_auth()`` method is used to explicitly set a default signing method. This method accepts one of four strings ``"api_key"``, ``"resource_principal"``, ``instance_principal`` or ``security_token``.
 
 The ``~/.oci/config`` configuration allow for multiple configurations to be stored in the same file. The ``set_auth()`` method takes is ``oci_config_location`` parameter that specifies the location of the configuration, and the default is ``"~/.oci/config"``. Each configuration is called a profile, and the default profile is ``DEFAULT``. The ``set_auth()`` method takes in a parameter ``profile``. It specifies which profile in the ``~/.oci/config`` configuration file to use. In this context, the ``profile`` parameter is only used when API keys are being used. If no value for ``profile`` is specified, then the ``DEFAULT`` profile section is used.
 
@@ -97,6 +112,7 @@ The ``~/.oci/config`` configuration allow for multiple configurations to be stor
   
   ads.set_auth("resource_principal")  # default signer is set to resource principal authentication
   ads.set_auth("instance_principal")  # default signer is set to instance principal authentication
+  ads.set_auth("security_token")  # default signer is set to security token authentication
 
   singer = oci.auth.signers.ResourcePrincipalsFederationSigner()
   ads.set_auth(config={}, singer=signer) # default signer is set to ResourcePrincipalsFederationSigner
@@ -122,9 +138,12 @@ Additional signers may be provided by running ``set_auth()`` with ``signer`` or
   oc.OCIClientFactory(**auth).object_storage
 
   # Example 3: Create Object Storage client with timeout set to 6000 using API Key authentication.
-  auth = authutil.api_keys(oci_config="/home/datascience/.oci/config", profile="TEST", kwargs={"timeout": 6000})
+  auth = authutil.api_keys(oci_config="/home/datascience/.oci/config", profile="TEST", client_kwargs={"timeout": 6000})
   oc.OCIClientFactory(**auth).object_storage
 
+  # Example 4: Create Object Storage client with timeout set to 6000 using security token authentication.
+  auth = authutil.security_token(oci_config="/home/datascience/.oci/config", profile="test_session", client_kwargs={"timeout": 6000})
+  oc.OCIClientFactory(**auth).object_storage
 
 In the this example, the default authentication uses API keys specified with the ``set_auth`` method. However, since the ``os_auth`` is specified to use resource principals, the notebook session uses the resource principal to access OCI Object Store.
 
@@ -144,11 +163,14 @@ More signers can be created using the ``create_signer()`` method. With the ``aut
   # Example 1. Create signer that uses instance principals
   auth = ads.auth.create_signer("instance_principal")
 
-  # Example 2. Provide a ResourcePrincipalsFederationSigner object
+  # Example 2. Create signer that uses security token
+  auth = ads.auth.create_signer("security_token", profile="test_session")
+
+  # Example 3. Provide a ResourcePrincipalsFederationSigner object
   singer = oci.auth.signers.ResourcePrincipalsFederationSigner()
   auth = ads.auth.create_signer(config={}, singer=signer)
 
-  # Example 3. Create signer that uses instance principals with log requests enabled
+  # Example 4. Create signer that uses instance principals with log requests enabled
   signer_callable = oci.auth.signers.InstancePrincipalsSecurityTokenSigner
   signer_kwargs = dict(log_requests=True) # will log the request url and response data when retrieving
   auth = ads.auth.create_signer(signer_callable=signer_callable, signer_kwargs=signer_kwargs)