Updated pr.

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
@@ -32,6 +34,11 @@
     "key_file", 
     "region"
 ]
+SECURITY_TOKEN_LEFT_TIME = 600
+
+
+class TokenExpiredError(Exception):  # pragma: no cover
+    pass
 
 
 class AuthType(str, metaclass=ExtendedEnumMeta):
@@ -819,21 +826,65 @@ def create_signer(self) -> Dict:
                     f"Parameter `{parameter}` must be provided for using `security_token` authentication."
                 )
 
-        if not self.oci_config:
-            os.system(f'oci session refresh --profile {self.oci_key_profile}')
-
         return {
             "config": configuration,
             "signer": oci.auth.signers.SecurityTokenSigner(
-                token=self._read_security_token_file(configuration.get("security_token_file")),
+                token=self._validate_and_refresh_token(configuration.get("security_token_file")),
                 private_key=oci.signer.load_private_key_from_file(configuration.get("key_file")),
                 generic_headers=configuration.get("generic_headers", SECURITY_TOKEN_GENERIC_HEADERS),
                 body_headers=configuration.get("body_headers", SECURITY_TOKEN_BODY_HEADERS)
             ),
             "client_kwargs": self.client_kwargs,
         }
-    
+
+    def _validate_and_refresh_token(self, security_token_file: str) -> str:
+        """Validates and refreshes security token.
+
+        Parameters
+        ----------
+        security_token_file: str
+            Path to security token file.
+
+        Returns
+        -------
+        str:
+            Security token string.
+        """
+        security_token = self._read_security_token_file(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 TokenExpiredError(
+                "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_now - time_expired < SECURITY_TOKEN_LEFT_TIME:
+            if self.oci_config_location == DEFAULT_LOCATION and self.oci_key_profile:
+                os.system(f'oci session refresh --profile {self.oci_key_profile}')
+                security_token = self._read_security_token_file(security_token_file)
+        
+        date_time = datetime.fromtimestamp(time_expired).strftime("%Y-%m-%d %H:%M:%S")
+        logger.info(f"Session is valid until {date_time}.")
+        return security_token
+
     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:

tests/unitary/default_setup/auth/test_auth.py

@@ -12,13 +12,16 @@
     EphemeralResourcePrincipalSigner,
 )
 from oci.auth.signers.security_token_signer import SecurityTokenSigner
+from oci.config import DEFAULT_LOCATION
 import ads
 from ads import set_auth
 from ads.common.utils import (
     oci_key_profile,
     oci_config_location,
 )
 from ads.common.auth import (
+    SecurityToken,
+    TokenExpiredError,
     api_keys,
     resource_principal,
     security_token,
@@ -124,97 +127,6 @@ def test_set_auth_with_key_content(self, mock_load_private_key, mock_validate_co
         assert signer["signer"] != None
         set_auth()
 
-    @mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
-    @mock.patch("oci.signer.load_private_key_from_file")
-    @mock.patch("ads.common.auth.SecurityToken._read_security_token_file")
-    def test_security_token_from_config(
-        self,
-        mock_read_security_token_file, 
-        mock_load_private_key_from_file,
-        mock_security_token_signer
-    ):
-        config = {
-            "fingerprint": "test_fingerprint",
-            "tenancy": "test_tenancy",
-            "region": "us-ashburn-1",
-            "key_file": "test_key_file",
-            "generic_headers": [1,2,3],
-            "body_headers": [4,5,6]
-        }
-
-        with pytest.raises(
-            ValueError,
-            match="Parameter `security_token_file` must be provided for using `security_token` authentication."
-        ):
-            signer = security_token(
-                oci_config=config,
-                client_kwargs={"test_client_key":"test_client_value"}
-            )
-
-        config["security_token_file"] = "test_security_token"
-        mock_security_token_signer.return_value = None
-        signer = security_token(
-            oci_config=config,
-            client_kwargs={"test_client_key":"test_client_value"}
-        )
-
-        mock_read_security_token_file.assert_called_with("test_security_token")
-        mock_load_private_key_from_file.assert_called_with("test_key_file")
-        assert signer["client_kwargs"] == {"test_client_key": "test_client_value"}
-        assert "additional_user_agent" in signer["config"]
-        assert signer["config"]["fingerprint"] == "test_fingerprint"
-        assert signer["config"]["tenancy"] == "test_tenancy"
-        assert signer["config"]["region"] == "us-ashburn-1"
-        assert signer["config"]["security_token_file"] == "test_security_token"
-        assert signer["config"]["key_file"] == "test_key_file"
-        assert isinstance(signer["signer"], SecurityTokenSigner)
-
-    @mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
-    @mock.patch("oci.signer.load_private_key_from_file")
-    @mock.patch("builtins.open")
-    @mock.patch("os.path.isfile")
-    @mock.patch("os.system")
-    @mock.patch("oci.config.from_file")
-    def test_security_token_from_file(
-        self,
-        mock_from_file,
-        mock_system,
-        mock_isfile,
-        mock_open,
-        mock_load_private_key_from_file,
-        mock_security_token_signer
-    ):
-        mock_from_file.return_value = {
-            "fingerprint": "test_fingerprint",
-            "tenancy": "test_tenancy",
-            "region": "us-ashburn-1",
-            "key_file": "test_key_file",
-            "security_token_file": "test_security_token"
-        }
-        mock_isfile.return_value = True
-        mock_security_token_signer.return_value = None
-        signer = security_token(
-            oci_config="test_config_location",
-            profile="test_key_profile",
-            client_kwargs={"test_client_key":"test_client_value"}
-        )
-
-        mock_from_file.assert_called_with("test_config_location", "test_key_profile")
-        mock_system.assert_called_with("oci session refresh --profile test_key_profile")
-        mock_isfile.assert_called_with("test_security_token")
-        mock_open.assert_called()
-        mock_load_private_key_from_file.assert_called_with("test_key_file")
-        mock_security_token_signer.assert_called()
-
-        assert signer["client_kwargs"] == {"test_client_key": "test_client_value"}
-        assert "additional_user_agent" in signer["config"]
-        assert signer["config"]["fingerprint"] == "test_fingerprint"
-        assert signer["config"]["tenancy"] == "test_tenancy"
-        assert signer["config"]["region"] == "us-ashburn-1"
-        assert signer["config"]["security_token_file"] == "test_security_token"
-        assert signer["config"]["key_file"] == "test_key_file"
-        assert isinstance(signer["signer"], SecurityTokenSigner)
-
 
 class TestOCIMixin(TestCase):
     def tearDown(self) -> None:
@@ -621,3 +533,110 @@ def test_with_set_auth_returns_error(self):
         with pytest.raises(ValueError):
             with AuthContext(auth="not_correct_auth_type"):
                 pass
+
+
+class TestSecurityToken(TestCase):
+
+    @mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
+    @mock.patch("oci.signer.load_private_key_from_file")
+    @mock.patch("ads.common.auth.SecurityToken._validate_and_refresh_token")
+    def test_security_token(
+        self,
+        mock_validate_and_refresh_token, 
+        mock_load_private_key_from_file,
+        mock_security_token_signer
+    ):
+        config = {
+            "fingerprint": "test_fingerprint",
+            "tenancy": "test_tenancy",
+            "region": "us-ashburn-1",
+            "key_file": "test_key_file",
+            "generic_headers": [1,2,3],
+            "body_headers": [4,5,6]
+        }
+
+        with pytest.raises(
+            ValueError,
+            match="Parameter `security_token_file` must be provided for using `security_token` authentication."
+        ):
+            signer = security_token(
+                oci_config=config,
+                client_kwargs={"test_client_key":"test_client_value"}
+            )
+
+        config["security_token_file"] = "test_security_token"
+        mock_security_token_signer.return_value = None
+        signer = security_token(
+            oci_config=config,
+            client_kwargs={"test_client_key":"test_client_value"}
+        )
+
+        mock_validate_and_refresh_token.assert_called_with("test_security_token")
+        mock_load_private_key_from_file.assert_called_with("test_key_file")
+        assert signer["client_kwargs"] == {"test_client_key": "test_client_value"}
+        assert "additional_user_agent" in signer["config"]
+        assert signer["config"]["fingerprint"] == "test_fingerprint"
+        assert signer["config"]["tenancy"] == "test_tenancy"
+        assert signer["config"]["region"] == "us-ashburn-1"
+        assert signer["config"]["security_token_file"] == "test_security_token"
+        assert signer["config"]["key_file"] == "test_key_file"
+        assert isinstance(signer["signer"], SecurityTokenSigner)
+
+    @mock.patch("os.system")
+    @mock.patch("oci.auth.security_token_container.SecurityTokenContainer.get_jwt")
+    @mock.patch("time.time")
+    @mock.patch("oci.auth.security_token_container.SecurityTokenContainer.valid")
+    @mock.patch("oci.auth.security_token_container.SecurityTokenContainer.__init__")
+    @mock.patch("ads.common.auth.SecurityToken._read_security_token_file")
+    def test_validate_and_refresh_token(
+        self, 
+        mock_read_security_token_file, 
+        mock_security_token_container,
+        mock_valid,
+        mock_time,
+        mock_get_jwt,
+        mock_system
+    ):
+        security_token = SecurityToken(
+            args={
+                "oci_config_location": DEFAULT_LOCATION,
+                "oci_key_profile": "test_profile"
+            }
+        )
+        mock_security_token_container.return_value = None
+        
+        mock_valid.return_value = False
+        with pytest.raises(
+            TokenExpiredError,
+            match="Security token has expired. Call `oci session authenticate` to generate new session."
+        ):
+            security_token._validate_and_refresh_token("test_security_token")
+        
+        
+        mock_valid.return_value = True
+        mock_time.return_value = 1
+        mock_get_jwt.return_value = {"exp" : 1}
+        
+        security_token._validate_and_refresh_token("test_security_token")
+        
+        mock_read_security_token_file.assert_called_with("test_security_token")
+        mock_security_token_container.assert_called()
+        mock_time.assert_called()
+        mock_get_jwt.assert_called()
+        mock_system.assert_called_with("oci session refresh --profile test_profile")
+
+    @mock.patch("builtins.open")
+    @mock.patch("os.path.isfile")
+    def test_read_security_token_file(self, mock_isfile, mock_open):
+        security_token = SecurityToken(args={})
+
+        mock_isfile.return_value = False
+        with pytest.raises(
+            ValueError,
+            match="Invalid `security_token_file`. Specify a valid path."
+        ):
+            security_token._read_security_token_file("test_security_token")
+
+        mock_isfile.return_value = True
+        security_token._read_security_token_file("test_security_token")
+        mock_open.assert_called()