From 36a3dd96ccf5449190f32fbe15c6fe48fe038b37 Mon Sep 17 00:00:00 2001 From: Pepe Barbe Date: Mon, 14 Apr 2025 18:07:01 -0500 Subject: [PATCH 1/5] feat: option to assume role to get credentials When fetching credentials after obtaining the boto3 session, we check if the configuration has an optional `assume_role` field, which indicates that role needs to be assumed to obtain the credentials. If found, the role is assumed and used to obtain the credentails. If not the original boto3 session is used to obtain the credentials. --- .gitignore | 1 + README.md | 4 ++ keyrings/codeartifact.py | 33 +++++++++++--- tests/config/single_with_role.cfg | 2 + tests/test_backend.py | 75 ++++++++++++++++++++++++------- tests/test_config.py | 25 ++++------- 6 files changed, 99 insertions(+), 41 deletions(-) create mode 100644 tests/config/single_with_role.cfg diff --git a/.gitignore b/.gitignore index 59da84b..c5d348d 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,4 @@ MANIFEST # Environments .venv venv/ +.idea/ diff --git a/README.md b/README.md index e9be3a8..984d5c6 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ Available options: - `token_duration`: Validity period (in seconds) for retieved authorization tokens. - `aws_access_key_id`: Use a specific AWS access key to authenticate with AWS. - `aws_secret_access_key`: Use a specific AWS secret access key to authenticate with AWS. + - `assume_role`: Role ARN to assume with the current profile name to get the CodeArtifact credentials. For more explanation of these options see the [AWS CLI documentation](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html). @@ -56,6 +57,9 @@ profile_name=default aws_access_key_id=xxxxxxxxx aws_secret_access_key=xxxxxxxxx +# Assume the following role to obtain the credentials +assume_role=arn:aws:iam::xxxxxxxxx:role/xxxxxxxxx + ``` ### Multiple Section Configuration (EXPERIMENTAL) diff --git a/keyrings/codeartifact.py b/keyrings/codeartifact.py index 1586e0e..f012602 100644 --- a/keyrings/codeartifact.py +++ b/keyrings/codeartifact.py @@ -156,13 +156,7 @@ def get_password(self, service, username): name=repository_name, ) - # Create session with any supplied configuration. - session = boto3.Session( - region_name=region, - profile_name=config.get("profile_name"), - aws_access_key_id=config.get("aws_access_key_id"), - aws_secret_access_key=config.get("aws_secret_access_key"), - ) + session = self.get_boto_session(region=region, config=config) # Create a CodeArtifact client for this repository's region. client = session.client("codeartifact", region_name=region) @@ -193,3 +187,28 @@ def set_password(self, service, username, password): def delete_password(self, service, username): # Defer deleting a password to the next backend raise NotImplementedError() + + @staticmethod + def get_boto_session(*, region, config): + # Create session with any supplied configuration. + session = boto3.Session( + region_name=region, + profile_name=config.get("profile_name"), + aws_access_key_id=config.get("aws_access_key_id"), + aws_secret_access_key=config.get("aws_secret_access_key"), + ) + + if "assume_role" in config: + sts_client = session.client("sts") + assumed_role_object = sts_client.assume_role( + RoleArn=config["assume_role"], + RoleSessionName="KeyRingsCodeArtifact", + ) + credentials = assumed_role_object["Credentials"] + return boto3.Session( + aws_access_key_id=credentials["AccessKeyId"], + aws_secret_access_key=credentials["SecretAccessKey"], + aws_session_token=credentials["SessionToken"], + ) + + return session diff --git a/tests/config/single_with_role.cfg b/tests/config/single_with_role.cfg new file mode 100644 index 0000000..7ebae22 --- /dev/null +++ b/tests/config/single_with_role.cfg @@ -0,0 +1,2 @@ +[codeartifact] +assume_role=arn:aws:iam::000000000000:role/some-role diff --git a/tests/test_backend.py b/tests/test_backend.py index c423966..51f9157 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -4,20 +4,28 @@ import keyring +from pathlib import Path from io import StringIO from urllib.parse import urlunparse from botocore.client import BaseClient from datetime import datetime, timedelta from keyrings.codeartifact import CodeArtifactBackend +CONFIG_DIR = Path(__file__).parent / "config" + @pytest.fixture -def backend(): +def backend(request): # Find the system-wide keyring. original = keyring.get_keyring() + if hasattr(request, "param") and "config_file" in request.param: + config_file = CONFIG_DIR / request.param["config_file"] + else: + config_file = StringIO() + # Use our keyring backend with an empty configuration. - backend = CodeArtifactBackend(config_file=StringIO()) + backend = CodeArtifactBackend(config_file=config_file) keyring.set_keyring(backend) yield backend @@ -67,26 +75,59 @@ def test_get_credential_invalid_path(backend, service): assert not keyring.get_credential(service, None) +def check_codeartifact_api_call(client, *args, **kwargs): + # We should only ever call GetAuthorizationToken + assert args[0] == "GetAuthorizationToken" + + # We should only ever supply these parameters. + assert args[1]["domain"] == "domain" + assert args[1]["domainOwner"] == "000000000000" + assert args[1]["durationSeconds"] == 3600 + + tzinfo = datetime.now().astimezone().tzinfo + current_time = datetime.now(tz=tzinfo) + + # Compute the expiration based on the current timestamp. + expiration = timedelta(seconds=args[1]["durationSeconds"]) + + return { + "authorizationToken": "TOKEN", + "expiration": current_time + expiration, + } + + def test_get_credential_supported_host(backend, monkeypatch): - def _make_api_call(client, *args, **kwargs): - # We should only ever call GetAuthorizationToken - assert args[0] == "GetAuthorizationToken" + monkeypatch.setattr(BaseClient, "_make_api_call", check_codeartifact_api_call) + url = codeartifact_pypi_url("domain", "000000000000", "region", "name") + credentials = backend.get_credential(url, None) - # We should only ever supply these parameters. - assert args[1]["domain"] == "domain" - assert args[1]["domainOwner"] == "000000000000" - assert args[1]["durationSeconds"] == 3600 + assert credentials.username == "aws" + assert credentials.password == "TOKEN" - tzinfo = datetime.now().astimezone().tzinfo - current_time = datetime.now(tz=tzinfo) - # Compute the expiration based on the current timestamp. - expiration = timedelta(seconds=args[1]["durationSeconds"]) +@pytest.mark.parametrize("backend", [{"config_file": "single_with_role.cfg"}], indirect=True) +def test_get_credential_for_assumed_role(backend, monkeypatch): + assumed_role = False - return { - "authorizationToken": "TOKEN", - "expiration": current_time + expiration, - } + def _make_api_call(client, *args, **kwargs): + nonlocal assumed_role + if not assumed_role: + # We should only ever call GetAuthorizationToken + assert args[0] == "AssumeRole" + + # We should only ever supply these parameters. + assert args[1]["RoleArn"] == "arn:aws:iam::000000000000:role/some-role" + assert args[1]["RoleSessionName"] == "KeyRingsCodeArtifact" + assumed_role = True + return { + "Credentials": { + "AccessKeyId": "", + "SecretAccessKey": "", + "SessionToken": "", + } + } + else: + return check_codeartifact_api_call(client, *args, **kwargs) monkeypatch.setattr(BaseClient, "_make_api_call", _make_api_call) url = codeartifact_pypi_url("domain", "000000000000", "region", "name") diff --git a/tests/test_config.py b/tests/test_config.py index a6f68bd..5594d1a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,21 +1,12 @@ # test_config.py -- config parsing tests - import pytest from io import StringIO -from os.path import dirname, join +from pathlib import Path from keyrings.codeartifact import CodeArtifactKeyringConfig - -@pytest.fixture -def config_file(): - working_directory = dirname(__file__) - - def _config_file(path): - return join(working_directory, "config", path) - - return _config_file +CONFIG_DIR = Path(__file__).parent / "config" @pytest.mark.parametrize( @@ -26,8 +17,8 @@ def _config_file(path): ("domain", "00000000", "ca-central-1", "repository"), ], ) -def test_parse_single_section_only(config_file, parameters): - config = CodeArtifactKeyringConfig(config_file("single_section.cfg")) +def test_parse_single_section_only(parameters): + config = CodeArtifactKeyringConfig(CONFIG_DIR / "single_section.cfg") # A single section has only one configuration. values = config.lookup(*parameters) @@ -89,8 +80,8 @@ def test_bogus_config_returns_empty_configuration(config_data): ), ], ) -def test_multiple_sections_with_defaults(config_file, query, expected): - path = config_file("multiple_sections_with_default.cfg") +def test_multiple_sections_with_defaults(query, expected): + path = CONFIG_DIR / "multiple_sections_with_default.cfg" config = CodeArtifactKeyringConfig(path) values = config.lookup(**query) @@ -115,8 +106,8 @@ def test_multiple_sections_with_defaults(config_file, query, expected): ), ], ) -def test_multiple_sections_no_defaults(config_file, query, expected): - path = config_file("multiple_sections_no_default.cfg") +def test_multiple_sections_no_defaults(query, expected): + path = CONFIG_DIR / "multiple_sections_no_default.cfg" config = CodeArtifactKeyringConfig(path) values = config.lookup(**query) From bed7137a8bb2d5c64061d0c161a998c7957176be Mon Sep 17 00:00:00 2001 From: Pepe Barbe Date: Thu, 17 Apr 2025 14:45:15 -0500 Subject: [PATCH 2/5] Address comments: * Run black. * Make role assumption logic more explicit. * Make role session name optional. --- README.md | 2 ++ keyrings/codeartifact.py | 16 ++++++++-------- tests/config/single_with_role.cfg | 1 + tests/test_backend.py | 11 +++++++---- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 984d5c6..0439f6e 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,8 @@ Available options: - `aws_access_key_id`: Use a specific AWS access key to authenticate with AWS. - `aws_secret_access_key`: Use a specific AWS secret access key to authenticate with AWS. - `assume_role`: Role ARN to assume with the current profile name to get the CodeArtifact credentials. + - `assume_role_session_name`: Name to attache to attach for the role session. If not specified, a name will be + selected by AWS SDK. For more explanation of these options see the [AWS CLI documentation](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html). diff --git a/keyrings/codeartifact.py b/keyrings/codeartifact.py index f012602..518154c 100644 --- a/keyrings/codeartifact.py +++ b/keyrings/codeartifact.py @@ -190,6 +190,8 @@ def delete_password(self, service, username): @staticmethod def get_boto_session(*, region, config): + should_assume_role = config.get("assume_role") + # Create session with any supplied configuration. session = boto3.Session( region_name=region, @@ -198,17 +200,15 @@ def get_boto_session(*, region, config): aws_secret_access_key=config.get("aws_secret_access_key"), ) - if "assume_role" in config: - sts_client = session.client("sts") - assumed_role_object = sts_client.assume_role( + if should_assume_role is not None: + assumed_role = session.client("sts").assume_role( RoleArn=config["assume_role"], - RoleSessionName="KeyRingsCodeArtifact", + RoleSessionName=config.get("assume_role_session_name"), ) - credentials = assumed_role_object["Credentials"] return boto3.Session( - aws_access_key_id=credentials["AccessKeyId"], - aws_secret_access_key=credentials["SecretAccessKey"], - aws_session_token=credentials["SessionToken"], + aws_access_key_id=assumed_role["Credentials"]["AccessKeyId"], + aws_secret_access_key=assumed_role["Credentials"]["SecretAccessKey"], + aws_session_token=assumed_role["Credentials"]["SessionToken"], ) return session diff --git a/tests/config/single_with_role.cfg b/tests/config/single_with_role.cfg index 7ebae22..471b3d2 100644 --- a/tests/config/single_with_role.cfg +++ b/tests/config/single_with_role.cfg @@ -1,2 +1,3 @@ [codeartifact] assume_role=arn:aws:iam::000000000000:role/some-role +assume_role_session_name=SomeSessionName diff --git a/tests/test_backend.py b/tests/test_backend.py index 51f9157..341c0d1 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -105,7 +105,9 @@ def test_get_credential_supported_host(backend, monkeypatch): assert credentials.password == "TOKEN" -@pytest.mark.parametrize("backend", [{"config_file": "single_with_role.cfg"}], indirect=True) +@pytest.mark.parametrize( + "backend", [{"config_file": "single_with_role.cfg"}], indirect=True +) def test_get_credential_for_assumed_role(backend, monkeypatch): assumed_role = False @@ -117,14 +119,14 @@ def _make_api_call(client, *args, **kwargs): # We should only ever supply these parameters. assert args[1]["RoleArn"] == "arn:aws:iam::000000000000:role/some-role" - assert args[1]["RoleSessionName"] == "KeyRingsCodeArtifact" + assert args[1]["RoleSessionName"] == "SomeSessionName" assumed_role = True return { - "Credentials": { + "Credentials": { "AccessKeyId": "", "SecretAccessKey": "", "SessionToken": "", - } + } } else: return check_codeartifact_api_call(client, *args, **kwargs) @@ -133,5 +135,6 @@ def _make_api_call(client, *args, **kwargs): url = codeartifact_pypi_url("domain", "000000000000", "region", "name") credentials = backend.get_credential(url, None) + assert assumed_role assert credentials.username == "aws" assert credentials.password == "TOKEN" From 22b1a172f83a70838cf4f51e285042d43d4cfe31 Mon Sep 17 00:00:00 2001 From: Pepe Barbe Date: Thu, 17 Apr 2025 15:40:05 -0500 Subject: [PATCH 3/5] Parametrize credential for supported host --- requirements-test.txt | 1 + tests/config/single_with_role.cfg | 3 - tests/test_backend.py | 154 +++++++++++++++++------------- 3 files changed, 88 insertions(+), 70 deletions(-) delete mode 100644 tests/config/single_with_role.cfg diff --git a/requirements-test.txt b/requirements-test.txt index c68f3a0..23383e4 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -1,2 +1,3 @@ pytest >= 6 pytest-cov +pytest-mock diff --git a/tests/config/single_with_role.cfg b/tests/config/single_with_role.cfg deleted file mode 100644 index 471b3d2..0000000 --- a/tests/config/single_with_role.cfg +++ /dev/null @@ -1,3 +0,0 @@ -[codeartifact] -assume_role=arn:aws:iam::000000000000:role/some-role -assume_role_session_name=SomeSessionName diff --git a/tests/test_backend.py b/tests/test_backend.py index 341c0d1..02ed317 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -4,28 +4,30 @@ import keyring -from pathlib import Path from io import StringIO from urllib.parse import urlunparse from botocore.client import BaseClient from datetime import datetime, timedelta -from keyrings.codeartifact import CodeArtifactBackend +from keyrings.codeartifact import CodeArtifactBackend, CodeArtifactKeyringConfig -CONFIG_DIR = Path(__file__).parent / "config" + +@pytest.fixture +def mocked_keyring_config(mocker): + mock_config_instance = mocker.create_autospec( + CodeArtifactKeyringConfig, spec_set=True + ) + mock_config = mocker.patch("keyrings.codeartifact.CodeArtifactKeyringConfig") + mock_config.return_value = mock_config_instance + return mock_config_instance @pytest.fixture -def backend(request): +def backend(mocked_keyring_config): # Find the system-wide keyring. original = keyring.get_keyring() - if hasattr(request, "param") and "config_file" in request.param: - config_file = CONFIG_DIR / request.param["config_file"] - else: - config_file = StringIO() - # Use our keyring backend with an empty configuration. - backend = CodeArtifactBackend(config_file=config_file) + backend = CodeArtifactBackend(config_file=StringIO()) keyring.set_keyring(backend) yield backend @@ -41,6 +43,55 @@ def codeartifact_pypi_url(domain, owner, region, name): return codeartifact_url(domain, owner, region, f"/pypi/{name}/") +def make_check_codeartifact_api_call(*, config, domain, domain_owner): + assumed_role = False + assume_role = config.get("assume_role") + assume_session_name = config.get("assume_session_name") + should_assume_role = assume_role is not None + + def _make_api_call(client, *args, **kwargs): + nonlocal assumed_role + if should_assume_role and not assumed_role: + # We should only ever call GetAuthorizationToken + assert args[0] == "AssumeRole" + + # We should only ever supply these parameters. + assert args[1]["RoleArn"] == assume_role + if assume_session_name is not None: + assert args[1]["RoleSessionName"] == assume_session_name + assumed_role = True + return { + "Credentials": { + "AccessKeyId": "", + "SecretAccessKey": "", + "SessionToken": "", + } + } + else: + assert assumed_role == should_assume_role + + # We should only ever call GetAuthorizationToken + assert args[0] == "GetAuthorizationToken" + + # We should only ever supply these parameters. + assert args[1]["domain"] == domain + assert args[1]["domainOwner"] == domain_owner + assert args[1]["durationSeconds"] == 3600 + + tzinfo = datetime.now().astimezone().tzinfo + current_time = datetime.now(tz=tzinfo) + + # Compute the expiration based on the current timestamp. + expiration = timedelta(seconds=args[1]["durationSeconds"]) + + return { + "authorizationToken": "TOKEN", + "expiration": current_time + expiration, + } + + return _make_api_call + + def test_set_password_raises(backend): with pytest.raises(NotImplementedError): keyring.set_password("service", "username", "password") @@ -75,66 +126,35 @@ def test_get_credential_invalid_path(backend, service): assert not keyring.get_credential(service, None) -def check_codeartifact_api_call(client, *args, **kwargs): - # We should only ever call GetAuthorizationToken - assert args[0] == "GetAuthorizationToken" - - # We should only ever supply these parameters. - assert args[1]["domain"] == "domain" - assert args[1]["domainOwner"] == "000000000000" - assert args[1]["durationSeconds"] == 3600 - - tzinfo = datetime.now().astimezone().tzinfo - current_time = datetime.now(tz=tzinfo) - - # Compute the expiration based on the current timestamp. - expiration = timedelta(seconds=args[1]["durationSeconds"]) - - return { - "authorizationToken": "TOKEN", - "expiration": current_time + expiration, - } - - -def test_get_credential_supported_host(backend, monkeypatch): - monkeypatch.setattr(BaseClient, "_make_api_call", check_codeartifact_api_call) - url = codeartifact_pypi_url("domain", "000000000000", "region", "name") - credentials = backend.get_credential(url, None) - - assert credentials.username == "aws" - assert credentials.password == "TOKEN" - - @pytest.mark.parametrize( - "backend", [{"config_file": "single_with_role.cfg"}], indirect=True + ["config"], + [ + ({},), + ( + { + "assume_role": "arn:aws:iam::000000000000:role/some-role", + "assume_role_session_name": "SomeSessionName", + }, + ), + ], ) -def test_get_credential_for_assumed_role(backend, monkeypatch): - assumed_role = False - - def _make_api_call(client, *args, **kwargs): - nonlocal assumed_role - if not assumed_role: - # We should only ever call GetAuthorizationToken - assert args[0] == "AssumeRole" - - # We should only ever supply these parameters. - assert args[1]["RoleArn"] == "arn:aws:iam::000000000000:role/some-role" - assert args[1]["RoleSessionName"] == "SomeSessionName" - assumed_role = True - return { - "Credentials": { - "AccessKeyId": "", - "SecretAccessKey": "", - "SessionToken": "", - } - } - else: - return check_codeartifact_api_call(client, *args, **kwargs) - - monkeypatch.setattr(BaseClient, "_make_api_call", _make_api_call) - url = codeartifact_pypi_url("domain", "000000000000", "region", "name") +def test_get_credential_supported_host( + backend, config, mocked_keyring_config, monkeypatch +): + domain = "domain" + domain_owner = "000000000000" + + monkeypatch.setattr( + BaseClient, + "_make_api_call", + make_check_codeartifact_api_call( + config=config, domain=domain, domain_owner=domain_owner + ), + ) + mocked_keyring_config.lookup.return_value = config + + url = codeartifact_pypi_url(domain, domain_owner, "region", "name") credentials = backend.get_credential(url, None) - assert assumed_role assert credentials.username == "aws" assert credentials.password == "TOKEN" From 89310a984aca9cdfefa8586bc80ce26ce95e591d Mon Sep 17 00:00:00 2001 From: Pepe Barbe Date: Tue, 22 Apr 2025 16:45:34 -0500 Subject: [PATCH 4/5] fix: settings role session name --- keyrings/codeartifact.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/keyrings/codeartifact.py b/keyrings/codeartifact.py index 518154c..48eb8b1 100644 --- a/keyrings/codeartifact.py +++ b/keyrings/codeartifact.py @@ -201,9 +201,12 @@ def get_boto_session(*, region, config): ) if should_assume_role is not None: + kwargs = {} + if "assume_role_session_name" in config: + kwargs["RoleSessionName"] = config["assume_role_session_name"] assumed_role = session.client("sts").assume_role( RoleArn=config["assume_role"], - RoleSessionName=config.get("assume_role_session_name"), + **kwargs, ) return boto3.Session( aws_access_key_id=assumed_role["Credentials"]["AccessKeyId"], From abfbd235b28130c3071c41f29e196f9518ca769e Mon Sep 17 00:00:00 2001 From: Pepe Barbe Date: Tue, 22 Apr 2025 16:48:23 -0500 Subject: [PATCH 5/5] provide default role assume session name --- keyrings/codeartifact.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/keyrings/codeartifact.py b/keyrings/codeartifact.py index 48eb8b1..1dcb345 100644 --- a/keyrings/codeartifact.py +++ b/keyrings/codeartifact.py @@ -201,12 +201,11 @@ def get_boto_session(*, region, config): ) if should_assume_role is not None: - kwargs = {} - if "assume_role_session_name" in config: - kwargs["RoleSessionName"] = config["assume_role_session_name"] assumed_role = session.client("sts").assume_role( RoleArn=config["assume_role"], - **kwargs, + RoleSessionName=config.get( + "assume_role_session_name", "KeyRingsCodeArtifact" + ), ) return boto3.Session( aws_access_key_id=assumed_role["Credentials"]["AccessKeyId"],