diff --git a/CHANGELOG.md b/CHANGELOG.md index 64e58ec..acee6b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Change Log +## [1.2.1](https://github.com/ome9ax/target-s3-jsonl/tree/1.2.1) (2022-07-13) +### What's Changed +* Added optional config parameter `role_arn`, which allows assuming additional roles. ## [1.2.0](https://github.com/ome9ax/target-s3-jsonl/tree/1.2.0) (2022-04-11) ### What's Changed diff --git a/README.md b/README.md index 623ae23..8466c29 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,7 @@ Full list of options in `config.json`: | aws_session_token | String | | AWS Session token. If not provided, `AWS_SESSION_TOKEN` environment variable will be used. | | encryption_type | String | | (Default: 'none') The type of encryption to use. Current supported options are: 'none' and 'KMS'. | | encryption_key | String | | A reference to the encryption key to use for data encryption. For KMS encryption, this should be the name of the KMS encryption key ID (e.g. '1234abcd-1234-1234-1234-1234abcd1234'). This field is ignored if 'encryption_type' is none or blank. | +| role_arn | String | | The ARN of the role to assume | ## Test Install the tools diff --git a/config.sample.json b/config.sample.json index c598259..9f45183 100644 --- a/config.sample.json +++ b/config.sample.json @@ -4,5 +4,6 @@ "s3_bucket": "BUCKET", "s3_key_prefix": "SOME-PREFIX/", "compression": "gzip", - "naming_convention": "{stream}-{timestamp}.jsonl" + "naming_convention": "{stream}-{timestamp}.jsonl", + "role_arn": "arn:aws:iam::000000000000:role/my_custom_role" } diff --git a/setup.cfg b/setup.cfg index 40672d0..ed6428e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -33,7 +33,7 @@ target_s3_jsonl = logging.conf [options.extras_require] test = pytest-cov - moto[s3] + moto[s3,sts] lint = flake8 dist = wheel diff --git a/target_s3_jsonl/__init__.py b/target_s3_jsonl/__init__.py index f67d4c8..2d4709f 100644 --- a/target_s3_jsonl/__init__.py +++ b/target_s3_jsonl/__init__.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -__version__ = '1.2.0' +__version__ = '1.2.1' import argparse import gzip @@ -238,6 +238,7 @@ def config_file(config_path): 'aws_session_token', 'aws_endpoint_url', 'aws_profile', + 'role_arn', 's3_bucket', 's3_key_prefix', 'encryption_type', diff --git a/target_s3_jsonl/s3.py b/target_s3_jsonl/s3.py index d4c5956..672e6a0 100644 --- a/target_s3_jsonl/s3.py +++ b/target_s3_jsonl/s3.py @@ -1,5 +1,7 @@ #!/usr/bin/env python3 import os +import re +import time import backoff import boto3 from botocore.exceptions import ClientError @@ -32,6 +34,9 @@ def create_client(config): aws_session_token = config.get('aws_session_token') or os.environ.get('AWS_SESSION_TOKEN') aws_profile = config.get('aws_profile') or os.environ.get('AWS_PROFILE') aws_endpoint_url = config.get('aws_endpoint_url') + role_arn = config.get('role_arn') + + endpoint_params = {'endpoint_url': aws_endpoint_url} if aws_endpoint_url else {} # AWS credentials based authentication if aws_access_key_id and aws_secret_access_key: @@ -43,10 +48,20 @@ def create_client(config): else: aws_session = boto3.session.Session(profile_name=aws_profile) - if aws_endpoint_url: - return aws_session.client('s3', endpoint_url=aws_endpoint_url) - else: - return aws_session.client('s3') + # config specifies a particular role to assume + # we create a session & s3-client with this role + if role_arn: + role_name = role_arn.split('/', 1)[1] + sts = aws_session.client('sts', **endpoint_params) + resp = sts.assume_role(RoleArn=role_arn, RoleSessionName=f'role-name={role_name}-profile={aws_profile}') + credentials = { + "aws_access_key_id": resp["Credentials"]["AccessKeyId"], + "aws_secret_access_key": resp["Credentials"]["SecretAccessKey"], + "aws_session_token": resp["Credentials"]["SessionToken"], + } + aws_session = boto3.Session(**credentials) + LOGGER.info(f"Creating s3 client with role {role_name}") + return aws_session.client('s3', **endpoint_params) # pylint: disable=too-many-arguments diff --git a/tests/resources/config_assume_role.json b/tests/resources/config_assume_role.json new file mode 100644 index 0000000..c07b07a --- /dev/null +++ b/tests/resources/config_assume_role.json @@ -0,0 +1,13 @@ +{ + "add_metadata_columns": false, + "aws_access_key_id": "ACCESS-KEY", + "aws_secret_access_key": "SECRET", + "s3_bucket": "BUCKET", + "temp_dir": "tests/output", + "memory_buffer": 2000000, + "compression": "none", + "timezone_offset": 0, + "naming_convention": "{stream}-{timestamp:%Y%m%dT%H%M%S}.jsonl", + "role_arn": "arn:aws:iam::123456789012:role/TestAssumeRole" + +} diff --git a/tests/test_s3.py b/tests/test_s3.py index 5cca495..d58857b 100644 --- a/tests/test_s3.py +++ b/tests/test_s3.py @@ -3,12 +3,12 @@ from copy import deepcopy from pathlib import Path import json - -# Third party imports import os -from pytest import fixture, raises +import re + import boto3 -from moto import mock_s3 +from moto import mock_s3, mock_sts +from pytest import fixture, raises # Package imports from target_s3_jsonl.s3 import create_client, upload_file, log_backoff_attempt @@ -22,6 +22,11 @@ def config(): return json.load(config_file) +@fixture +def config_assume_role(): + with open(Path('tests', 'resources', 'config_assume_role.json'), 'r', encoding='utf-8') as f: + return json.load(f) + @fixture(scope='module') def aws_credentials(): """Mocked AWS Credentials for moto.""" @@ -36,9 +41,16 @@ def test_log_backoff_attempt(caplog): '''TEST : simple upload_files call''' log_backoff_attempt({'tries': 99}) + pat = r'INFO root:s3.py:\d{2} Error detected communicating with Amazon, triggering backoff: 99 try\n' + assert re.match(pat, caplog.text) - assert caplog.text == 'INFO root:s3.py:22 Error detected communicating with Amazon, triggering backoff: 99 try\n' +@mock_sts +@mock_s3 +def test_create_client_with_assumed_role(config_assume_role, caplog): + """Assert client is created with assumed role when role_arn is specified""" + client = create_client(config_assume_role) + assert caplog.text.endswith('Creating s3 client with role TestAssumeRole\n') @mock_s3 def test_create_client(aws_credentials, config):