feat: Add file caching (#990)

* Add file cache

* feat: add output file cache support

Author: Chuan Ren

google/auth/pluggable.py

@@ -41,6 +41,7 @@
 import json
 import os
 import subprocess
+import time
 
 from google.auth import _helpers
 from google.auth import exceptions
@@ -163,6 +164,17 @@ def retrieve_subject_token(self, request):
                 "Executables need to be explicitly allowed (set GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES to '1') to run."
             )
 
+        # Check output file
+        if self._credential_source_executable_output_file is not None:
+            try:
+                with open(self._credential_source_executable_output_file) as output_file:
+                    response = json.load(output_file)
+                    subject_token = self._parse_subject_token(response)
+            except:
+                pass
+            else:
+                return subject_token
+
         # Inject env vars
         original_audience = os.getenv("GOOGLE_EXTERNAL_ACCOUNT_AUDIENCE")
         os.environ["GOOGLE_EXTERNAL_ACCOUNT_AUDIENCE"] = self._audience
@@ -208,23 +220,8 @@ def retrieve_subject_token(self, request):
         else:
             data = result.stdout.decode('utf-8')
             response = json.loads(data)
-            if not response['success']:
-                raise exceptions.RefreshError(
-                    "Executable returned unsuccessful response: {}.".format(response)
-                )
-            elif response['version'] > EXECUTABLE_SUPPORTED_MAX_VERSION:
-                raise exceptions.RefreshError(
-                    "Executable returned unsupported version {}.".format(response['version'])
-                )
-            elif response["token_type"] == "urn:ietf:params:oauth:token-type:jwt" or response["token_type"] == "urn:ietf:params:oauth:token-type:id_token": # OIDC
-                return response["id_token"]
-            elif response["token_type"] == "urn:ietf:params:oauth:token-type:saml2": # SAML
-                return response["saml_response"]
-            else:
-                raise exceptions.RefreshError(
-                    "Executable returned unsupported token type."
-                ) 
-
+            return self._parse_subject_token(response)
+            
     @classmethod
     def from_info(cls, info, **kwargs):
         """Creates a Pluggable Credentials instance from parsed external account info.
@@ -271,3 +268,25 @@ def from_file(cls, filename, **kwargs):
         with io.open(filename, "r", encoding="utf-8") as json_file:
             data = json.load(json_file)
             return cls.from_info(data, **kwargs)
+
+    def _parse_subject_token(self, response):
+        if not response['success']:
+                raise exceptions.RefreshError(
+                    "Executable returned unsuccessful response: {}.".format(response)
+                )
+        elif response['version'] > EXECUTABLE_SUPPORTED_MAX_VERSION:
+            raise exceptions.RefreshError(
+                "Executable returned unsupported version {}.".format(response['version'])
+            )
+        elif response['expiration_time'] < time.time():
+            raise exceptions.RefreshError(
+                "The token returned by the executable is expired."
+            )
+        elif response["token_type"] == "urn:ietf:params:oauth:token-type:jwt" or response["token_type"] == "urn:ietf:params:oauth:token-type:id_token": # OIDC
+            return response["id_token"]
+        elif response["token_type"] == "urn:ietf:params:oauth:token-type:saml2": # SAML
+            return response["saml_response"]
+        else:
+            raise exceptions.RefreshError(
+                "Executable returned unsupported token type."
+            ) 

tests/test_pluggable.py

@@ -49,10 +49,11 @@
 
 class TestCredentials(object):
     CREDENTIAL_SOURCE_EXECUTABLE_COMMAND = "/fake/external/excutable --arg1=value1 --arg2=value2"
+    CREDENTIAL_SOURCE_EXECUTABLE_OUTPUT_FILE = "fake_output_file"
     CREDENTIAL_SOURCE_EXECUTABLE = {
         "command": CREDENTIAL_SOURCE_EXECUTABLE_COMMAND,
         "timeout_millis": 5000,
-        "output_file": "/fake/output/file"
+        "output_file": CREDENTIAL_SOURCE_EXECUTABLE_OUTPUT_FILE
     }
     CREDENTIAL_SOURCE = {
         "executable": CREDENTIAL_SOURCE_EXECUTABLE
@@ -78,7 +79,7 @@ class TestCredentials(object):
         "success": True,
         "token_type": "urn:ietf:params:oauth:token-type:saml2",
         "saml_response": EXECUTABLE_SAML_TOKEN,
-        "expiration_time": 1620433341
+        "expiration_time": 9999999999
     }
     EXECUTABLE_FAILED_RESPONSE = {
         "version": 1,
@@ -488,7 +489,20 @@ def test_retrieve_subject_token_failed(self, fp):
             subject_token = credentials.retrieve_subject_token(None)
 
         assert excinfo.match(r"Executable returned unsuccessful response")
-        
+
+    @mock.patch.dict(os.environ, {"GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES": "0"})
+    def test_retrieve_subject_token_not_allowd(self, fp):
+        fp.register(self.CREDENTIAL_SOURCE_EXECUTABLE_COMMAND.split(), stdout=json.dumps(self.EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_ID_TOKEN))
+    
+        credentials = self.make_pluggable(
+            credential_source=self.CREDENTIAL_SOURCE
+        )
+
+        with pytest.raises(ValueError) as excinfo:
+            subject_token = credentials.retrieve_subject_token(None)
+
+        assert excinfo.match(r"Executables need to be explicitly allowed")
+
     @mock.patch.dict(os.environ, {"GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES": "1"})
     def test_retrieve_subject_token_invalid_version(self, fp):
         EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_VERSION_2 = {
@@ -498,7 +512,7 @@ def test_retrieve_subject_token_invalid_version(self, fp):
             "id_token": self.EXECUTABLE_OIDC_TOKEN,
             "expiration_time": 9999999999
         }
-            
+
         fp.register(self.CREDENTIAL_SOURCE_EXECUTABLE_COMMAND.split(), stdout=json.dumps(EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_VERSION_2))
 
         credentials = self.make_pluggable(
@@ -508,4 +522,41 @@ def test_retrieve_subject_token_invalid_version(self, fp):
         with pytest.raises(exceptions.RefreshError) as excinfo:
             subject_token = credentials.retrieve_subject_token(None)
 
-        assert excinfo.match(r"Executable returned unsupported version")
+        assert excinfo.match(r"Executable returned unsupported version")
+
+    @mock.patch.dict(os.environ, {"GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES": "1"})
+    def test_retrieve_subject_token_expired_token(self, fp):
+        EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_EXPIRED= {
+            "version": 1,
+            "success": True,
+            "token_type": "urn:ietf:params:oauth:token-type:id_token",
+            "id_token": self.EXECUTABLE_OIDC_TOKEN,
+            "expiration_time": 0
+        }
+            
+        fp.register(self.CREDENTIAL_SOURCE_EXECUTABLE_COMMAND.split(), stdout=json.dumps(EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_EXPIRED))
+    
+        credentials = self.make_pluggable(
+            credential_source=self.CREDENTIAL_SOURCE
+        )
+
+        with pytest.raises(exceptions.RefreshError) as excinfo:
+            subject_token = credentials.retrieve_subject_token(None)
+
+        assert excinfo.match(r"The token returned by the executable is expired")
+        
+    @mock.patch.dict(os.environ, {"GOOGLE_EXTERNAL_ACCOUNT_ALLOW_EXECUTABLES": "1"})
+    def test_retrieve_subject_token_file_cache(self, fp):
+        with open(self.CREDENTIAL_SOURCE_EXECUTABLE_OUTPUT_FILE, 'w') as output_file:
+            json.dump(self.EXECUTABLE_SUCCESSFUL_OIDC_RESPONSE_ID_TOKEN, output_file)
+    
+        credentials = self.make_pluggable(
+            credential_source=self.CREDENTIAL_SOURCE
+        )
+
+        subject_token = credentials.retrieve_subject_token(None)
+
+        assert subject_token == self.EXECUTABLE_OIDC_TOKEN
+        
+        if os.path.exists(self.CREDENTIAL_SOURCE_EXECUTABLE_OUTPUT_FILE):
+            os.remove(self.CREDENTIAL_SOURCE_EXECUTABLE_OUTPUT_FILE)