# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2021
# SPDX-License-Identifier: Apache-2.0

import argparse
import uuid
import boto3
import os
import sys

sys.path.append(os.path.join(os.path.dirname(__file__), '../../libs'))
import deploy_utils

'''
This utility creates an IAM role with necessary permissions for the CookieFactory workspace in Grafana
'''

def parse_args():
    parser = argparse.ArgumentParser(
        description='Creates a role for the AWS IoT TwinMaker workspace for this CookieFactory sample to be used in Grafana.')
    parser.add_argument('--workspace-id',
                        help='workspace to be used on a Grafana dashboard',
                        required=True, default='CookieFactory')
    parser.add_argument('--account-id',
                        help='accountId for workspace',
                        required=True)
    parser.add_argument('--auth-provider',
                        help="(optional) If you are using Amazon Managed Grafana, enter the Workspace IAM Role ARN in this field.",
                        required=False, default=None)
    parser.add_argument('--region',
                        help="(optional) AWS region you are creating the sample in. Defaults to 'us-east-1'",
                        required=False, default='us-east-1')
    parser.add_argument('--profile',
                        help="(optional) AWS profile to access your account with. See your configured profiles with "
                             "`~/.aws/credentials`. Defaults to 'None'",
                        required=False, default=None)
    parser.add_argument('--endpoint-url', required=False, default=None, help='AWS IoT TwinMaker service endpoint')
    return parser.parse_args()


def main():
    args = parse_args()
    region = args.region
    profile = args.profile
    workspaceId = args.workspace_id
    endpoint_url = args.endpoint_url
    auth_provider = args.auth_provider
    accountId = args.account_id

    session = boto3.session.Session(profile_name=profile)
    iam = session.client(service_name='iam', region_name=region)
    sts = session.client(service_name='sts', region_name=region)
    iottwinmaker = session.client(service_name='iottwinmaker', endpoint_url=endpoint_url, region_name=region)
    identifier = uuid.uuid4().hex

    userArn = sts.get_caller_identity()['Arn']
    if auth_provider == None:
      auth_provider = userArn

    ws = deploy_utils.WorkspaceUtils(
        workspace_id=workspaceId,
        region_name=region,
        endpoint_url=endpoint_url,
        profile=profile)

    role_name = f"{'Cloud9-' if is_c9() else ''}IoTTwinMakerDashboardRole-{identifier[:8]}"
    create_role_response = iam.create_role(
        RoleName=role_name,
        Path="/",
        AssumeRolePolicyDocument=f"""{{
                "Version": "2012-10-17",
                "Statement": [
                    {{
                        "Effect": "Allow",
                        "Principal": {{
                            "AWS": "{auth_provider}"
                        }},
                        "Action": "sts:AssumeRole"
                    }}
                ]
            }}
        """,
        Description='Role generated by AWS IoT TwinMaker Getting Started for a Grafana dashboard',
    )
    role_arn = create_role_response['Role']['Arn']
    print(f"Created IAM role: {role_arn}")

    workspace = iottwinmaker.get_workspace(workspaceId=workspaceId)
    s3BucketArn = workspace.get('s3Location')
    workspaceArn = workspace.get('arn')
    create_policy_response = iam.create_policy(
        PolicyName=f"{workspaceId}DashboardPolicy-{identifier}",
        PolicyDocument=f"""{{
          "Version": "2012-10-17",
          "Statement": [
            {{
              "Effect": "Allow",
              "Action": ["s3:GetObject"],
              "Resource": [
                "{s3BucketArn}", 
                "{s3BucketArn}/*"
              ]
            }},
            {{
              "Effect": "Allow",
              "Action": ["iottwinmaker:Get*", "iottwinmaker:List*"],
              "Resource": [
                "{workspaceArn}",
                "{workspaceArn}/*"
              ]
            }},
            {{
              "Effect": "Allow",
              "Action": "iottwinmaker:ListWorkspaces",
              "Resource": "*"
            }},
            {{
              "Effect": "Allow",
              "Action": [
                "kinesisvideo:GetDataEndpoint",
                "kinesisvideo:GetHLSStreamingSessionURL"
              ],
              "Resource": [
                "arn:aws:kinesisvideo:{region}:{accountId}:stream/cookiefactory_mixerroom_camera_01/*", 
                "arn:aws:kinesisvideo:{region}:{accountId}:stream/cookiefactory_mixerroom_camera_02/*"
              ]
            }}
          ]
	      }}"""
    )

    policy_arn = create_policy_response['Policy']['Arn']
    print(f"Created the dashboard IAM policy: {policy_arn}")
    iam.attach_role_policy(
        RoleName=role_name,
        PolicyArn=policy_arn
    )
    print(f"Attached {workspaceId}DashboardPolicy-{identifier} to dashboard role.")
    
    ws.store_sample_metadata("samples_content_dashboard_role_name", role_name)

    print(f"Please use this Role ARN when configuring the IoT TwinMaker datasource in Grafana:\n"
          f"\n  {role_arn}\n")

# check if we're running in cloud9, if so then created roles need to be prefixed with "Cloud9-"
# https://docs.aws.amazon.com/cloud9/latest/user-guide/security-iam.html#auth-and-access-control-temporary-managed-credentials-supported
def is_c9():
    return os.getenv("C9_PROJECT") is not None

if __name__ == '__main__':
    main()