From 31bec1955bb324e0951572271835f27165955d7a Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Fri, 10 Jan 2025 14:00:58 +0000 Subject: [PATCH] fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! General improvements --- docker-image-worker/export.py | 5 +++-- docker-image-worker/s3_utils.py | 39 +++++++++++++++++++++++++-------- internal/shared/jobs.go | 12 ++++++---- 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/docker-image-worker/export.py b/docker-image-worker/export.py index 6cdd6c9..e4e671c 100755 --- a/docker-image-worker/export.py +++ b/docker-image-worker/export.py @@ -31,7 +31,7 @@ def log_error_details(e): logger.error(f"Non-AWS Error: {str(e)}") @retry(stop=stop_after_attempt(5), wait=wait_fixed(5)) -def transfer_file(source, destination, use_role=False, role_name=None, aws_access_key_id=None, aws_secret_access_key=None, endpoint_url=None, region=None): +def transfer_file(source, destination, use_role=False, role_name=None, use_current_role=False, aws_access_key_id=None, aws_secret_access_key=None, endpoint_url=None, region=None): """ Transfer a file from a local source to either a local destination or an S3 bucket """ @@ -43,7 +43,7 @@ def transfer_file(source, destination, use_role=False, role_name=None, aws_acces # Uploading to S3 try: logger.info(f"Attempting to upload {source} to {destination}") - s3_client = get_s3_client(use_role, role_name, aws_access_key_id, aws_secret_access_key, endpoint_url, region) + s3_client = get_s3_client(use_role, role_name, use_current_role, aws_access_key_id, aws_secret_access_key, endpoint_url, region) bucket, s3_key = parse_s3_path(destination) try: @@ -92,6 +92,7 @@ if __name__ == "__main__": args.destination, args.use_role, args.role_name, + args.use_current_role, args.aws_access_key_id, args.aws_secret_access_key, args.endpoint_url, diff --git a/docker-image-worker/s3_utils.py b/docker-image-worker/s3_utils.py index 7de2f64..206bc24 100644 --- a/docker-image-worker/s3_utils.py +++ b/docker-image-worker/s3_utils.py @@ -4,7 +4,7 @@ from botocore.exceptions import ClientError import os import logging -def get_s3_client(use_role=False, role_name=None, aws_access_key_id=None, aws_secret_access_key=None, endpoint_url=None, region=None): +def get_s3_client(use_role=False, role_name=None, use_current_role=False, aws_access_key_id=None, aws_secret_access_key=None, endpoint_url=None, region=None): """ Create and return an S3 client based on the provided authentication method, endpoint, and region. """ @@ -48,9 +48,9 @@ def get_s3_client(use_role=False, role_name=None, aws_access_key_id=None, aws_se logger.info(f"Current identity: {identity['Arn']}") assumed_role_object = sts_client.assume_role( - RoleArn=f"arn:aws:iam::{boto3.client('sts').get_caller_identity()['Account']}:role/{role_name}", - RoleSessionName="AssumeRoleSession" - ) + RoleArn=f"arn:aws:iam::{boto3.client('sts').get_caller_identity()['Account']}:role/{role_name}", + RoleSessionName="AssumeRoleSession" + ) credentials = assumed_role_object['Credentials'] client_kwargs['aws_access_key_id'] = credentials['AccessKeyId'] client_kwargs['aws_secret_access_key'] = credentials['SecretAccessKey'] @@ -59,6 +59,19 @@ def get_s3_client(use_role=False, role_name=None, aws_access_key_id=None, aws_se except Exception as e: logger.error(f"Failed to assume role {role_name}: {str(e)}") raise + elif use_current_role: + # Use the current role (e.g., from Kubernetes service account) + logger.info("Using current role from environment") + try: + client = boto3.client('s3', **client_kwargs) + # Try to get caller identity to verify credentials + sts = boto3.client('sts') + identity = sts.get_caller_identity() + logger.info(f"Successfully authenticated using current role: {identity['Arn']}") + return client + except Exception as e: + logger.error(f"Failed to use current role: {str(e)}") + raise else: # Use default credentials (environment, instance profile, or pod service account) logger.info("Using default credential provider chain") @@ -86,8 +99,10 @@ def add_common_arguments(parser): """ Add common command-line arguments to an ArgumentParser object """ - parser.add_argument("--use_role", action="store_true", help="Use IAM role for authentication") - parser.add_argument("--role_name", help="The name of the IAM role to assume") + auth_group = parser.add_mutually_exclusive_group() + auth_group.add_argument("--use_role", action="store_true", help="Use IAM role for authentication") + auth_group.add_argument("--use_current_role", action="store_true", help="Use current AWS role (e.g. from Kubernetes service account)") + parser.add_argument("--role_name", help="The name of the IAM role to assume (only when --use_role is set)") parser.add_argument("--aws_access_key_id", help="AWS access key ID") parser.add_argument("--aws_secret_access_key", help="AWS secret access key") parser.add_argument("--endpoint_url", help="S3-compatible endpoint URL") @@ -99,9 +114,15 @@ def validate_args(args, parser): """ if args.destination.startswith('s3://'): # Check for conflicting auth methods - if args.use_role and args.role_name and (args.aws_access_key_id or args.aws_secret_access_key): - parser.error("When using a specific role (--role_name), access key and secret should not be specified.") + if args.use_role and not args.role_name: + parser.error("--role_name is required when using --use_role") + + if args.role_name and not args.use_role: + parser.error("--role_name can only be used with --use_role") + + if args.use_current_role and (args.aws_access_key_id or args.aws_secret_access_key): + parser.error("When using current role (--use_current_role), access key and secret should not be specified") # If using explicit credentials, require both key and secret if (args.aws_access_key_id or args.aws_secret_access_key) and not (args.aws_access_key_id and args.aws_secret_access_key): - parser.error("Both --aws_access_key_id and --aws_secret_access_key must be provided when using access key authentication.") + parser.error("Both --aws_access_key_id and --aws_secret_access_key must be provided when using access key authentication") diff --git a/internal/shared/jobs.go b/internal/shared/jobs.go index b209308..8cd4307 100644 --- a/internal/shared/jobs.go +++ b/internal/shared/jobs.go @@ -102,14 +102,18 @@ func CreateJob[T any](params JobParams, setupFunc func(T) []string) *batchv1.Job func SetupS3Params(s3Config raczylocomv1.ClusterImageStorageS3) []string { params := []string{} if s3Config.UseRole { - params = append(params, "--use_role") + if s3Config.RoleARN != "" { + // Use specific role + params = append(params, "--use_role") + params = append(params, fmt.Sprintf("--role_name='%s'", s3Config.RoleARN)) + } else { + // Use current role from Kubernetes service account + params = append(params, "--use_current_role") + } } else { params = append(params, fmt.Sprintf("--aws_access_key_id='%s'", s3Config.AccessKey)) params = append(params, fmt.Sprintf("--aws_secret_access_key='%s'", s3Config.SecretKey)) } - if s3Config.RoleARN != "" { - params = append(params, fmt.Sprintf("--role_name='%s'", s3Config.RoleARN)) - } if s3Config.Endpoint != "" { params = append(params, fmt.Sprintf("--endpoint_url='%s'", s3Config.Endpoint)) }