Updated pr.

Author: lu-ohai

ads/common/auth.py

@@ -6,13 +6,17 @@
 
 import copy
 from datetime import datetime
+import json
 import os
 from dataclasses import dataclass
 import time
 from typing import Any, Callable, Dict, Optional, Union
 
+import requests
+
 import ads.telemetry
 import oci
+from oci_cli import cli_util
 from ads.common import logger
 from ads.common.decorator.deprecate import deprecated
 from ads.common.extended_enum import ExtendedEnumMeta
@@ -37,7 +41,7 @@
 SECURITY_TOKEN_LEFT_TIME = 600
 
 
-class TokenExpiredError(Exception):  # pragma: no cover
+class SecurityTokenError(Exception):  # pragma: no cover
     pass
 
 
@@ -825,52 +829,106 @@ def create_signer(self) -> Dict:
                 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._validate_and_refresh_token(configuration.get("security_token_file")),
-                private_key=oci.signer.load_private_key_from_file(configuration.get("key_file")),
+                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", 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:
+    def _validate_and_refresh_token(self, configuration: Dict[str, Any]):
         """Validates and refreshes security token.
 
         Parameters
         ----------
-        security_token_file: str
-            Path to security token file.
-
-        Returns
-        -------
-        str:
-            Security token string.
+        configuration: Dict
+            Security token configuration.
         """
-        security_token = self._read_security_token_file(security_token_file)
+        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 TokenExpiredError(
+            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_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)
+            try:
+                self._refresh_security_token(configuration)
+            except Exception as ex:
+                logger.info("Failed to refresh security token. Error: {}".format(ex))
 
         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 _refresh_security_token(self, configuration: Dict[str, Any]):
+        """Refreshes security token. The logic is mainly taken reference from:
+        https://github.com/oracle/oci-cli/blob/9a0978344950d7b7c24a688892f24968dce20ad3/src/oci_cli/cli_session.py#L152
+
+        Parameters
+        ----------
+        configuration: Dict
+            Security token configuration.
+        """
+        expanded_security_token_location = os.path.expanduser(
+            configuration.get("security_token_file")
+        )
+
+        with open(expanded_security_token_location, 'r') as security_token_file:
+            token = security_token_file.read()
+
+        try:
+            private_key = oci.signer.load_private_key_from_file(
+                configuration.get("key_file"), configuration.get("pass_phrase")
+            )
+        except:
+            raise
+        auth = oci.auth.signers.SecurityTokenSigner(token, private_key)
+
+        refresh_url = "{endpoint}/v1/authentication/refresh".format(
+            endpoint=oci.regions.endpoint_for("auth", configuration.get("region"))
+        )
+        logger.info(f"Attempting to refresh token from {refresh_url}.")
+
+        response = requests.post(
+            refresh_url,
+            headers={
+                'content-type': 'application/json'
+            },
+            data=json.dumps({
+                'currentToken': token
+            }),
+            auth=auth
+        )
+
+        if response.status_code == 200:
+            refreshed_token = json.loads(response.content.decode('UTF-8'))['token']
+            with open(expanded_security_token_location, 'w') as security_token_file:
+                security_token_file.write(refreshed_token)
+            cli_util.apply_user_only_access_permissions(expanded_security_token_location)
+            logger.info("Successfully refreshed token")
+        elif response.status_code == 401:
+            raise SecurityTokenError(
+                "Security token has expired. Call `oci session authenticate` to generate new session."
+            )
+        else:
+            raise SecurityTokenError(
+                "Failed to refresh sesison. Error: {}".format(str(response.content.decode('UTF-8')))
+            )
 
     def _read_security_token_file(self, security_token_file: str) -> str:
         """Reads security token from file.

tests/unitary/default_setup/auth/test_auth.py

@@ -4,6 +4,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import os
+from mock import MagicMock
 import pytest
 from unittest import TestCase, mock
 
@@ -21,7 +22,7 @@
 )
 from ads.common.auth import (
     SecurityToken,
-    TokenExpiredError,
+    SecurityTokenError,
     api_keys,
     resource_principal,
     security_token,
@@ -539,10 +540,12 @@ 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._read_security_token_file")
     @mock.patch("ads.common.auth.SecurityToken._validate_and_refresh_token")
     def test_security_token(
         self,
         mock_validate_and_refresh_token, 
+        mock_read_security_token_file,
         mock_load_private_key_from_file,
         mock_security_token_signer
     ):
@@ -571,8 +574,9 @@ def test_security_token(
             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")
+        mock_validate_and_refresh_token.assert_called_with(config)
+        mock_read_security_token_file.assert_called_with("test_security_token")
+        mock_load_private_key_from_file.assert_called_with("test_key_file", None)
         assert signer["client_kwargs"] == {"test_client_key": "test_client_value"}
         assert "additional_user_agent" in signer["config"]
         assert signer["config"]["fingerprint"] == "test_fingerprint"
@@ -582,7 +586,7 @@ def test_security_token(
         assert signer["config"]["key_file"] == "test_key_file"
         assert isinstance(signer["signer"], SecurityTokenSigner)
 
-    @mock.patch("os.system")
+    @mock.patch("ads.common.auth.SecurityToken._refresh_security_token")
     @mock.patch("oci.auth.security_token_container.SecurityTokenContainer.get_jwt")
     @mock.patch("time.time")
     @mock.patch("oci.auth.security_token_container.SecurityTokenContainer.valid")
@@ -595,7 +599,7 @@ def test_validate_and_refresh_token(
         mock_valid,
         mock_time,
         mock_get_jwt,
-        mock_system
+        mock_refresh_security_token
     ):
         security_token = SecurityToken(
             args={
@@ -606,24 +610,94 @@ def test_validate_and_refresh_token(
         mock_security_token_container.return_value = None
 
         mock_valid.return_value = False
+        configuration = {
+            "fingerprint": "test_fingerprint",
+            "tenancy": "test_tenancy",
+            "region": "us-ashburn-1",
+            "key_file": "test_key_file",
+            "security_token_file": "test_security_token",
+            "generic_headers": [1,2,3],
+            "body_headers": [4,5,6]
+        }
         with pytest.raises(
-            TokenExpiredError,
+            SecurityTokenError,
             match="Security token has expired. Call `oci session authenticate` to generate new session."
         ):
-            security_token._validate_and_refresh_token("test_security_token")
+            security_token._validate_and_refresh_token(configuration)
 
 
         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")
+        security_token._validate_and_refresh_token(configuration)
 
         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_refresh_security_token.assert_called_with(configuration)
+
+    @mock.patch("oci_cli.cli_util.apply_user_only_access_permissions")
+    @mock.patch("json.loads")
+    @mock.patch("requests.post")
+    @mock.patch("json.dumps")
+    @mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
+    @mock.patch("oci.signer.load_private_key_from_file")
+    @mock.patch("builtins.open")
+    def test_refresh_security_token(
+        self,
+        mock_open,
+        mock_load_private_key_from_file,
+        mock_security_token_signer,
+        mock_dumps,
+        mock_post,
+        mock_loads,
+        mock_apply_user_only_access_permissions
+    ):
+        security_token = SecurityToken(args={})
+        configuration = {
+            "fingerprint": "test_fingerprint",
+            "tenancy": "test_tenancy",
+            "region": "us-ashburn-1",
+            "key_file": "test_key_file",
+            "security_token_file": "test_security_token",
+            "generic_headers": [1,2,3],
+            "body_headers": [4,5,6]
+        }
+        mock_security_token_signer.return_value = None
+        mock_loads.return_value = {
+            "token": "test_token"
+        }
+
+        response = MagicMock()
+        response.status_code = 401
+        mock_post.return_value = response
+        with pytest.raises(
+            SecurityTokenError,
+            match="Security token has expired. Call `oci session authenticate` to generate new session."
+        ):
+            security_token._refresh_security_token(configuration)
+
+        response.status_code = 500
+        mock_post.return_value = response
+        with pytest.raises(
+            SecurityTokenError,
+        ):
+            security_token._refresh_security_token(configuration)
+
+        response.status_code = 200
+        response.content = bytes("test_content", encoding='utf8')
+        mock_post.return_value = response
+        security_token._refresh_security_token(configuration)
+
+        mock_open.assert_called()
+        mock_load_private_key_from_file.assert_called_with("test_key_file", None)
+        mock_security_token_signer.assert_called()
+        mock_dumps.assert_called()
+        mock_post.assert_called()
+        mock_loads.assert_called()
+        mock_apply_user_only_access_permissions.assert_called()
 
     @mock.patch("builtins.open")
     @mock.patch("os.path.isfile")