Parameters:
  SharedServicesAccountID:
    Description: 'Account ID to host shared services'
    Type: String
  SharedS3Bucket:
    Description: 'S3 bucket name of shared artifacts'
    Type: String 
  SharedKMSKeyARN:
    Description: 'ARN of shared KMS key in A'
    Type: String

Resources:
  S3BucketAccessForCrossAccountRoleC:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          -
            Effect: Allow
            Action:
            - 's3:GetObject*'
            - 's3:PutObject'
            - 's3:PutObjectAcl'
            - 'codecommit:ListBranches'
            - 'codecommit:ListRepositories'
            Resource: !Sub 'arn:aws:s3:::${SharedS3Bucket}/*'
  KMSKeyAccessForC:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          -
            Effect: Allow
            Action:
            - 'kms:DescribeKey'
            - 'kms:GenerateDataKey*'
            - 'kms:Encrypt'
            - 'kms:ReEncrypt'
            - 'kms:Decrypt'
            Resource: !Ref 'SharedKMSKeyARN'
  PassRoleforC:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          -
            Effect: Allow
            Action:
            - 'iam:PassRole'
            Resource: 
            - !Sub 'arn:aws:iam::${AWS::AccountId}:role/${AWS::StackName}*'
            - 'arn:aws:codepipeline:us-east-1:557045202904:mvp-codepipeline'
  CAssumeARole:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          -
            Effect: Allow
            Action:
            - 'sts:AssumeRole'
            Resource: 
            - !Sub 'arn:aws:iam::${SharedServicesAccountID}:role/*'
  CrossAccountRoleC:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - 
            Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
                - sagemaker.amazonaws.com
            Action: sts:AssumeRole
          - 
            Effect: Allow
            Principal:
              AWS: 
                - !Sub 'arn:aws:iam::${SharedServicesAccountID}:root'
            Action: sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - !Ref S3BucketAccessForCrossAccountRoleC
        - !Ref KMSKeyAccessForC
        - !Ref PassRoleforC
        - !Ref CAssumeARole
        - 'arn:aws:iam::aws:policy/AWSCodePipelineCustomActionAccess'
        - 'arn:aws:iam::aws:policy/AWSLambda_FullAccess'
  DeployModelLambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt CrossAccountRoleC.Arn
      Code:
        ZipFile: 
            |
            import json
            import boto3
            import os
               
            def lambda_handler(event, context):
                # TODO implement
                print(event)
                
                print(event['CodePipeline.job']['data']['actionConfiguration']['configuration']['UserParameters'])
                
                STS_SESSION = get_sts_session(event, \
                event['CodePipeline.job']['accountId'], \
                event['CodePipeline.job']['data']['actionConfiguration']['configuration']['UserParameters'])
                iam = STS_SESSION.client('iam')
                codepipeline = STS_SESSION.client('codepipeline')
                codepipeline.put_job_success_result(jobId=event['CodePipeline.job']['id'])
                return 0

            def put_job_success(event):
                
                print("[SUCCESS]Endpoint Deployed")
                print(event['message'])
                code_pipeline.put_job_success_result(jobId=event['CodePipeline.job']['id'])

            def put_job_failure(event):  
                print('[ERROR]Putting job failure')
                print(event['message'])
                code_pipeline.put_job_success_result(jobId=event['CodePipeline.job']['id'])
                
            def get_sts_session(event, account, rolename):
                sts = boto3.client("sts")
                RoleArn = str("arn:aws:iam::" + account + ":role/" + rolename)
                print(RoleArn)
                response = sts.assume_role(
                    RoleArn=RoleArn,
                    RoleSessionName='SecurityManageAccountPermissions',
                    DurationSeconds=900)
                sts_session = boto3.Session(
                    aws_access_key_id=response['Credentials']['AccessKeyId'],
                    aws_secret_access_key=response['Credentials']['SecretAccessKey'],
                    aws_session_token=response['Credentials']['SessionToken'],
                    region_name=os.environ['AWS_REGION'],
                    botocore_session=None,
                    profile_name=None)
                return (sts_session)

      Runtime: python3.6
      Timeout: 500
      MemorySize: 1000