Skip to content

Commit b2020e7

Browse files
lu-ohailu-ohai
committed
Updated pr.
1 parent 22d0b1d commit b2020e7

File tree

2 files changed

+74
-43
lines changed

2 files changed

+74
-43
lines changed

‎ads/common/auth.py‎

Lines changed: 33 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,22 @@
1717
from oci.config import DEFAULT_LOCATION # "~/.oci/config"
1818
from oci.config import DEFAULT_PROFILE # "DEFAULT"
1919

20+
SECURITY_TOKEN_GENERIC_HEADERS = [
21+
"date",
22+
"(request-target)",
23+
"host"
24+
]
25+
SECURITY_TOKEN_BODY_HEADERS = [
26+
"content-length",
27+
"content-type",
28+
"x-content-sha256"
29+
]
30+
SECURITY_TOKEN_REQUIRED = [
31+
"security_token_file",
32+
"key_file",
33+
"region"
34+
]
35+
2036

2137
class AuthType(str, metaclass=ExtendedEnumMeta):
2238
API_KEY = "api_key"
@@ -144,7 +160,8 @@ def set_auth(
144160
>>> ads.set_auth("security_token") # Set security token authentication
145161
146162
>>> config = dict(
147-
... key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem
163+
... region=us-ashburn-1,
164+
... key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem,
148165
... security_token_file=~/.oci/sessions/DEFAULT/token
149166
... )
150167
>>> ads.set_auth("security_token", config=config) # Set security token authentication from provided config
@@ -400,7 +417,8 @@ def create_signer(
400417
>>> signer_kwargs = dict(log_requests=True) # will log the request url and response data when retrieving
401418
>>> auth = ads.auth.create_signer(signer_callable=signer_callable, signer_kwargs=signer_kwargs) # instance principals authentication dictionary created based on callable with kwargs parameters
402419
>>> config = dict(
403-
... key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem
420+
... region=us-ashburn-1,
421+
... key_file=~/.oci/sessions/DEFAULT/oci_api_key.pem,
404422
... security_token_file=~/.oci/sessions/DEFAULT/token
405423
... )
406424
>>> auth = ads.auth.create_signer(auth_type="security_token", config=config) # security token authentication created based on provided config
@@ -795,38 +813,29 @@ def create_signer(self) -> Dict:
795813

796814
logger.info(f"Using 'security_token' authentication.")
797815

798-
if "security_token_file" not in configuration and "security_token_content" not in configuration:
799-
raise ValueError(
800-
"Parameter `security_token_file` or `security_token_content` must be provided for using `security_token` authentication."
801-
)
802-
803-
if "key_file" not in configuration and "key_content" not in configuration:
804-
raise ValueError(
805-
"Parameter `key_file` or `key_content` must be provided for using `security_token` authentication."
806-
)
816+
for parameter in SECURITY_TOKEN_REQUIRED:
817+
if parameter not in configuration:
818+
raise ValueError(
819+
f"Parameter `{parameter}` must be provided for using `security_token` authentication."
820+
)
807821

808-
if "security_token_content" not in configuration and not self.oci_config:
809-
os.system(f'oci session refresh --profile {self.oci_key_profile or DEFAULT_PROFILE}')
822+
if not self.oci_config:
823+
os.system(f'oci session refresh --profile {self.oci_key_profile}')
810824

811825
return {
812826
"config": configuration,
813827
"signer": oci.auth.signers.SecurityTokenSigner(
814-
token=(
815-
configuration.get("security_token_content", None)
816-
or self._read_security_token_file(configuration.get("security_token_file"))
817-
),
818-
private_key=(
819-
oci.signer.load_private_key(configuration.get("key_content"))
820-
if configuration.get("key_content")
821-
else oci.signer.load_private_key_from_file(configuration.get("key_file"))
822-
),
823-
generic_headers=configuration.get("generic_headers"),
824-
body_headers=configuration.get("body_headers")
828+
token=self._read_security_token_file(configuration.get("security_token_file")),
829+
private_key=oci.signer.load_private_key_from_file(configuration.get("key_file")),
830+
generic_headers=configuration.get("generic_headers", SECURITY_TOKEN_GENERIC_HEADERS),
831+
body_headers=configuration.get("body_headers", SECURITY_TOKEN_BODY_HEADERS)
825832
),
826833
"client_kwargs": self.client_kwargs,
827834
}
828835

829836
def _read_security_token_file(self, security_token_file: str) -> str:
837+
if not os.path.isfile(security_token_file):
838+
raise ValueError("Invalid `security_token_file`. Specify a valid path.")
830839
try:
831840
token = None
832841
with open(security_token_file, 'r') as f:

‎tests/unitary/default_setup/auth/test_auth.py‎

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -125,44 +125,33 @@ def test_set_auth_with_key_content(self, mock_load_private_key, mock_validate_co
125125
set_auth()
126126

127127
@mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
128-
@mock.patch("oci.signer.load_private_key")
129128
@mock.patch("oci.signer.load_private_key_from_file")
130129
@mock.patch("ads.common.auth.SecurityToken._read_security_token_file")
131-
def test_security_token(
130+
def test_security_token_from_config(
132131
self,
133132
mock_read_security_token_file,
134133
mock_load_private_key_from_file,
135-
mock_load_private_key,
136134
mock_security_token_signer
137135
):
138136
config = {
139137
"fingerprint": "test_fingerprint",
140138
"tenancy": "test_tenancy",
141139
"region": "us-ashburn-1",
140+
"key_file": "test_key_file",
142141
"generic_headers": [1,2,3],
143142
"body_headers": [4,5,6]
144143
}
145144

146145
with pytest.raises(
147146
ValueError,
148-
match="Parameter `security_token_file` or `security_token_content` must be provided for using `security_token` authentication."
147+
match="Parameter `security_token_file` must be provided for using `security_token` authentication."
149148
):
150149
signer = security_token(
151150
oci_config=config,
152151
client_kwargs={"test_client_key":"test_client_value"}
153152
)
154153

155154
config["security_token_file"] = "test_security_token"
156-
with pytest.raises(
157-
ValueError,
158-
match="Parameter `key_file` or `key_content` must be provided for using `security_token` authentication."
159-
):
160-
signer = security_token(
161-
oci_config=config,
162-
client_kwargs={"test_client_key":"test_client_value"}
163-
)
164-
165-
config["key_file"] = "test_key_file"
166155
mock_security_token_signer.return_value = None
167156
signer = security_token(
168157
oci_config=config,
@@ -180,18 +169,51 @@ def test_security_token(
180169
assert signer["config"]["key_file"] == "test_key_file"
181170
assert isinstance(signer["signer"], SecurityTokenSigner)
182171

183-
config = {
172+
@mock.patch("oci.auth.signers.SecurityTokenSigner.__init__")
173+
@mock.patch("oci.signer.load_private_key_from_file")
174+
@mock.patch("builtins.open")
175+
@mock.patch("os.path.isfile")
176+
@mock.patch("os.system")
177+
@mock.patch("oci.config.from_file")
178+
def test_security_token_from_file(
179+
self,
180+
mock_from_file,
181+
mock_system,
182+
mock_isfile,
183+
mock_open,
184+
mock_load_private_key_from_file,
185+
mock_security_token_signer
186+
):
187+
mock_from_file.return_value = {
184188
"fingerprint": "test_fingerprint",
185189
"tenancy": "test_tenancy",
186190
"region": "us-ashburn-1",
187-
"security_token_content": "test_security_token_content",
188-
"key_content": "test_key_content"
191+
"key_file": "test_key_file",
192+
"security_token_file": "test_security_token"
189193
}
194+
mock_isfile.return_value = True
195+
mock_security_token_signer.return_value = None
190196
signer = security_token(
191-
oci_config=config,
197+
oci_config="test_config_location",
198+
profile="test_key_profile",
192199
client_kwargs={"test_client_key":"test_client_value"}
193200
)
194-
mock_load_private_key.assert_called_with("test_key_content")
201+
202+
mock_from_file.assert_called_with("test_config_location", "test_key_profile")
203+
mock_system.assert_called_with("oci session refresh --profile test_key_profile")
204+
mock_isfile.assert_called_with("test_security_token")
205+
mock_open.assert_called()
206+
mock_load_private_key_from_file.assert_called_with("test_key_file")
207+
mock_security_token_signer.assert_called()
208+
209+
assert signer["client_kwargs"] == {"test_client_key": "test_client_value"}
210+
assert "additional_user_agent" in signer["config"]
211+
assert signer["config"]["fingerprint"] == "test_fingerprint"
212+
assert signer["config"]["tenancy"] == "test_tenancy"
213+
assert signer["config"]["region"] == "us-ashburn-1"
214+
assert signer["config"]["security_token_file"] == "test_security_token"
215+
assert signer["config"]["key_file"] == "test_key_file"
216+
assert isinstance(signer["signer"], SecurityTokenSigner)
195217

196218

197219
class TestOCIMixin(TestCase):

0 commit comments

Comments
 (0)