Skip to content

Commit

Permalink
feat: Add file caching (#990)
Browse files Browse the repository at this point in the history
* Add file cache

* feat: add output file cache support
  • Loading branch information
Chuan Ren committed May 27, 2022
1 parent 6f80803 commit 31444d9
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 22 deletions.
53 changes: 36 additions & 17 deletions google/auth/pluggable.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
import json
import os
import subprocess
import time

from google.auth import _helpers
from google.auth import exceptions
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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."
)
61 changes: 56 additions & 5 deletions tests/test_pluggable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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 = {
Expand All @@ -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(
Expand All @@ -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)

0 comments on commit 31444d9

Please sign in to comment.