in migration/bring-your-own-role/byor.py [0:0]
def byor_main():
args = _parse_args()
session = boto3.Session()
if (args.region):
session = boto3.Session(region_name=args.region)
iam_client = session.client('iam')
datazone = session.client('datazone')
lakeformation = session.client('lakeformation')
sagemaker = session.client('sagemaker')
if args.endpoint:
datazone = session.client('datazone', endpoint_url=args.endpoint)
if args.command == ROLE_REPLACEMENT:
print(f"Use bring in Role: {args.bring_in_role_arn} as Project Role...")
# Get Project's Auto Generated Execution Role, there should be one role per project
project_role = _find_project_execution_role(args, iam_client, datazone)
# Get Execution Role's trust policy
project_role_trust_policy = project_role['Role']['AssumeRolePolicyDocument']
byor_role = iam_client.get_role(
RoleName=_get_role_name_from_arn(args.bring_in_role_arn),
)
environment_with_role_lists = _get_enviroments_with_role_from_project(datazone, args, project_role['Role']['Arn'])
# Replace Project Execution Role with BYOR Role
# Role is attached with environment, and one Project contains multiple environments, so
# we need to replace role for each environment within a project
for environment in environment_with_role_lists:
print(f"Will replace IAM role {environment.user_role_arn} attached to environment name: {environment.name}, id: {environment.id} with new role {args.bring_in_role_arn}...\n")
if args.execute:
try:
print(f"Disassociate role {environment.user_role_arn} from environment {environment.id} in progress... \n")
response = datazone.disassociate_environment_role(
domainIdentifier=args.domain_id,
environmentIdentifier=environment.id,
environmentRoleArn=environment.user_role_arn
)
print(f"Successfully disassociate role {environment.user_role_arn} from environment {environment.id}: {response} \n")
except ClientError as e:
if e.response['Error']['Code'] == 'ResourceNotFoundException':
print(f"Disassociate role {environment.user_role_arn} from environment {environment.id} failed: Role not found in environment, skip disassociate. \n")
else:
raise e
print(f"Associate role {args.bring_in_role_arn} to environment {environment.id} in progress... \n")
try:
response = datazone.associate_environment_role(
domainIdentifier=args.domain_id,
environmentIdentifier=environment.id,
environmentRoleArn=args.bring_in_role_arn
)
print(f"Associate role {args.bring_in_role_arn} to environment {environment.id} successfully: {response} \n")
except Exception as e:
# Associate environment role failed, re-associate with original role
print(f"Associate role {args.bring_in_role_arn} to environment {environment.id} failed: {e}, re-associate with original role {environment.user_role_arn}. But all subscriptions are lost, please recreate necessary subscriptions.\n")
response = datazone.associate_environment_role(
domainIdentifier=args.domain_id,
environmentIdentifier=environment.id,
environmentRoleArn=environment.user_role_arn
)
raise e
else:
print(f"Skipping disassociate and associate role operations, set --execute flag to True to do the actual update. environment {environment.name} still use {environment.user_role_arn} as its role.\n")
# Copy DataZone Subscriptions
if not environment.name == 'RedshiftServerless' and not environment.name == 'Redshift Serverless':
_copy_datazone_subscriptions(args.domain_id, environment.id, datazone, byor_role, args.execute)
# Copy LakeFormation Permissions and Opt-Ins
_copy_lakeformation_grants(lakeformation, environment.user_role_arn, args.bring_in_role_arn, args.execute, args.command)
_copy_lakeformation_opt_ins(lakeformation, environment.user_role_arn, args.bring_in_role_arn, args.execute)
# Get BYOR Role's trust policy
byor_role_trust_policy = byor_role['Role']['AssumeRolePolicyDocument']
# Combine trust policy and update BYOR Role's trust policy
new_trust_policy = _combine_trust_policy(project_role_trust_policy, byor_role_trust_policy)
_update_trust_policy(byor_role['Role']['RoleName'], new_trust_policy, iam_client, args.execute)
# Copy Project Execution Role's managed policies to BYOR Role
_copy_managed_policies_arn(project_role, byor_role, iam_client, args.execute)
# Copy Project Execution Role's inline policies to BYOR Role
_copy_inline_policies_arn(project_role, byor_role, iam_client, args.execute)
# Copy Project Execution Role's Tags to BYOR Role
_copy_tags(project_role['Role']['RoleName'], byor_role['Role']['RoleName'], iam_client, args.execute)
# Update associated EMR Instance Role's policies
emr_instance_role_policies = _find_emr_instance_role_policies(args, iam_client)
if emr_instance_role_policies is not None:
_replace_role_arn_in_policies(emr_instance_role_policies,
iam_client,
project_role['Role']['Arn'],
args.bring_in_role_arn,
args.execute)
# Replace SageMaker Domain Execution Role
sagemaker_domain_id = _find_sagemaker_domain_id(sagemaker, args)
if sagemaker_domain_id:
if args.force_update:
_stop_apps_under_domain(sagemaker, sagemaker_domain_id, args.execute)
else:
raise Exception(f"Updating SageMaker Domain without deleting running apps is failing this script execution. Set --force-update flag if you accept app deletion to ensure successful script execution.")
_update_domain_execution_role(sagemaker, sagemaker_domain_id, args.bring_in_role_arn, args.execute)
# Update LakeFormation Data lake locations resources with the new Role
_update_s3_lakeformation_registration(lakeformation, project_role['Role']['Arn'], args.bring_in_role_arn, args.execute)
# Create or update SMUS Provisioning Role's inline policy
_update_smus_provisioning_role(datazone, iam_client, args.domain_id, args.bring_in_role_arn, args.execute)
if args.execute:
print(f"Successfully replace Project {args.project_id} user role with your own role: {byor_role['Role']['Arn']}")
elif args.command == ROLE_ENHANCEMENT:
print(f"Enhance Project Role...")
# Get Project's Auto Generated Role
project_role = _find_project_execution_role(args, iam_client, datazone)
# Get Project Role's trust policy
project_role_trust_policy = project_role['Role']['AssumeRolePolicyDocument']
# Get BYOR Role's trust policy
byor_role = iam_client.get_role(
RoleName=_get_role_name_from_arn(args.bring_in_role_arn),
)
print(f"BYOR Role ARN: {args.bring_in_role_arn}\n")
byor_role_trust_policy = byor_role['Role']['AssumeRolePolicyDocument']
# Combine trust policy and update Project Role's trust policy
new_trust_policy = _combine_trust_policy(project_role_trust_policy, byor_role_trust_policy)
_update_trust_policy(project_role['Role']['RoleName'], new_trust_policy, iam_client, args.execute)
# Copy BYOR Role's managed policies to Project Role
_copy_managed_policies_arn(byor_role, project_role, iam_client, args.execute)
# Copy BYOR Role's inline policies to Project Role
_copy_inline_policies_arn(byor_role, project_role, iam_client, args.execute)
# Copy BYOR Role's Tags to Project Role
_copy_tags(byor_role['Role']['RoleName'], project_role['Role']['RoleName'], iam_client, args.execute)
# Copy LakeFormation Permissions and Opt-Ins
_copy_lakeformation_grants(lakeformation, args.bring_in_role_arn, project_role['Role']['Arn'], args.execute, args.command)
_copy_lakeformation_opt_ins(lakeformation, args.bring_in_role_arn, project_role['Role']['Arn'], args.execute)
if args.execute:
print(f"Successfully enhance project user role: {project_role['Role']['Arn']} referring to your own role: {byor_role['Role']['Arn']}")
else:
print(f"Invalid command. Expecting '{ROLE_REPLACEMENT}' or '{ROLE_ENHANCEMENT}'.")