---
AWSTemplateFormatVersion: 2010-09-09
Description: Mavvrik AWS Usage Account Onboarding v3 - Main Stack

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label: { "default": "User Configuration" }
        Parameters: [tenant, token, accessType]
      - Label: { "default": "System Configuration (Optional)" }
        Parameters: [parentStackId, parentStackOrgUnitIds]
    ParameterLabels:
      tenant: { default: Tenant }
      token: { default: Token (Optional) }
      accessType: { default: Access Type (Optional) }
      parentStackId: { default: Parent Stack (Optional) }
      parentStackOrgUnitIds: { default: Parent Stack's Organizational Unit IDs (Optional) }

Parameters:
  tenant:
    Type: String
    MinLength: 3
    Description: Mavvrik assigned tenant id. Mostly this will be auto populated if not, get it from Mavvrik Console.

  token:
    Type: String
    NoEcho: true
    Description: Not supposed to be changed manually. Leave empty if its not automatically populated

  accessType:
    Type: String
    Default: ""
    AllowedPattern: ^(|x-loginasadmin|x-loginasroot)$
    ConstraintDescription: must be empty, x-loginasadmin, or x-loginasroot
    Description: Not supposed to be changed manually. Leave empty if its not automatically populated

  parentStackId:
    Type: String
    Description: Not supposed to be changed manually. Leave empty if its not automatically populated

  parentStackOrgUnitIds:
    Type: CommaDelimitedList
    Description: Not supposed to be changed manually. Leave empty if its not automatically populated

Mappings:
  Mavvrik:
    AssumeRole:
      Principals:
        - arn:aws:iam::911403356698:user/global-user-for-assume-role
        - arn:aws:iam::911403356698:user/mavvrik-svc-user
        - arn:aws:iam::911403356698:role/mavvrik-svc-role

Conditions:
  TokenAvailable: !Not [!Equals [!Ref token, ""]]
  ParentAvailable: !Not [!Equals [!Ref parentStackId, ""]]

Resources:
  ResourceDiscoveryRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "usage-discovery-${tenant}"
      Path: /mavvrik/
      Description: Mavvrik > Read only access to discover resources
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action: "sts:AssumeRole"
            Principal: { AWS: !FindInMap [Mavvrik, AssumeRole, Principals] }
            Condition: { StringEquals: { "sts:ExternalId": !Ref tenant } }
          - Effect: Allow
            Action: "sts:TagSession"
            Principal: { AWS: !FindInMap [Mavvrik, AssumeRole, Principals] }
            Condition:
              StringEquals: { "aws:RequestTag/mavvrik_tenant": !Ref tenant }
              "ForAllValues:StringEquals": { "aws:TagKeys": ["mavvrik_tenant"] }
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/ReadOnlyAccess
        - arn:aws:iam::aws:policy/job-function/ViewOnlyAccess
        - arn:aws:iam::aws:policy/IAMReadOnlyAccess
        - arn:aws:iam::aws:policy/CloudWatchReadOnlyAccess
        - arn:aws:iam::aws:policy/ComputeOptimizerReadOnlyAccess
        - arn:aws:iam::aws:policy/AWSBillingReadOnlyAccess
      Policies:
        - PolicyName: AllowCommitmentRecommendations
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "ce:StartSavingsPlansPurchaseRecommendationGeneration"
                  - "ce:GetSavingsPlansPurchaseRecommendation"
                  - "ce:GetReservationPurchaseRecommendation"
                Resource: "*"
        - PolicyName: DenySSMParameterStoreAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Deny
                Action:
                  - "ssm:GetParameter"
                  - "ssm:GetParameters"
                  - "ssm:GetParameterHistory"
                  - "ssm:GetParametersByPath"
                  - "ssm:DescribeParameters"
                  - "ssm:*Parameter*"
                Resource: "*"
      Tags: [{ Key: "app", Value: "Mavvrik" }, { Key: "version", Value: "v3" }, { Key: "variant", Value: "v3" }]

  LambdaFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Path: /mavvrik/
      Description: Mavvrik > A role required to run a lambda function which makes an API calls to Mavvrik platform to onboard a account
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          Effect: Allow
          Action: sts:AssumeRole
          Principal: { Service: "lambda.amazonaws.com" }
      Policies:
        - PolicyName: AllowEnrollingForComputeOptimizer
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: "compute-optimizer:*EnrollmentStatus"
                Resource: "*"
              - Effect: Allow
                Action: "iam:CreateServiceLinkedRole"
                Resource: !Sub "arn:aws:iam::${AWS::AccountId}:role/*"
                Condition:
                  StringEquals:
                    iam:AWSServiceName: compute-optimizer.amazonaws.com
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Tags: [{ Key: "app", Value: "Mavvrik" }, { Key: "version", Value: "v3" }, { Key: "variant", Value: "v3" }]

  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Description: Mavvrik > function to onboard this account to Mavvrik
      Role: !GetAtt LambdaFunctionExecutionRole.Arn
      Runtime: python3.13
      Handler: index.handler
      Timeout: 600
      Code:
        ZipFile: |
          import json
          import random
          import traceback
          from time import sleep
          from typing import Callable
          from typing import TypedDict
          from http.client import HTTPSConnection, HTTPResponse
          import boto3
          import cfnresponse
          HTTP_SUCCESS_CODES = list(range(200, 211))
          HTTP_AUTH_FAILURE_CODES = [401, 403]
          class Payload(TypedDict):
              Host: str
              Token: str
              AccessType: str
              Tenant: str
              StackId: str
              RoleARN: str
              AccountId: str
              OrgUnitIds: list[str]
          def retry(func: Callable):
              def wrapper(*args, **kwargs):
                  err = None
                  for _ in range(0, 10):
                      try:
                          return func(*args, **kwargs)
                      except Exception as e:
                          print(f"request failed : {e}")
                          err = e
                          sleep(random.randint(1, 10))
                  raise err
              return wrapper
          class Onboarding:
              def __init__(self, payload: Payload):
                  self.payload = payload
              @property
              def _id(self):
                  return f"aws-{self.payload['AccountId']}"
              @property
              def url_prefix(self):
                  return f"/{self.payload['Tenant']}/usage-account"
              @property
              def headers(self):
                  headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.payload['Token']}"}
                  access_type = self.payload.get("AccessType", "")
                  if access_type and access_type.strip():
                      headers[access_type] = "true"
                  return headers
              @property
              def request_payload(self):
                  return {
                      "id": self._id,
                      "version": "v3",
                      "cfn_version": "v3",
                      "cfn_id": self.payload["StackId"],
                      "cfn_param_ous": self.payload["OrgUnitIds"] or [],
                      "provider_id": "aws",
                      "account_id": self.payload["AccountId"],
                      "rolearn": self.payload["RoleARN"],
                  }
              def request(self, method, url, data: dict = None) -> HTTPResponse:
                  print(f"calling {method} {url}")
                  connection = HTTPSConnection(self.payload["Host"])
                  connection.request(method, url, headers=self.headers, body=json.dumps(data) if data else None)
                  response = connection.getresponse()
                  print(f"status code : {response.status}")
                  return response
              def exists(self):
                  url = f"{self.url_prefix}/{self._id}"
                  response = self.request("GET", url)
                  if response.status in HTTP_SUCCESS_CODES:
                      data = json.loads(response.read().decode())
                      return data.get("status", False)
                  return False
              @retry
              def create_account(self):
                  method, url = "POST", self.url_prefix
                  response = self.request(method, url, self.request_payload)
                  if response.status in HTTP_AUTH_FAILURE_CODES:
                      print(f"token expired or unauthorized ({response.status}), skipping")
                      return
                  if response.status not in HTTP_SUCCESS_CODES:
                      raise Exception(f"{method} {url} {response.status} {response.read()}")
              @retry
              def update_account(self):
                  method, url = "PUT", f"{self.url_prefix}/{self._id}"
                  response = self.request(method, url, self.request_payload)
                  if response.status in HTTP_AUTH_FAILURE_CODES:
                      print(f"token expired or unauthorized ({response.status}), skipping")
                      return
                  if response.status not in HTTP_SUCCESS_CODES:
                      raise Exception(f"{method} {url} {response.status} {response.read()}")
              def onboard(self):
                  print("checking if account is already onboarded")
                  if self.exists():
                      print("account already exist, updating account")
                      return self.update_account()
                  print("account does not exist already, creating account")
                  self.create_account()
              @retry
              def delete_account(self):
                  method, url = "DELETE", f"{self.url_prefix}/{self._id}"
                  response = self.request(method, url)
                  if response.status not in HTTP_SUCCESS_CODES:
                      raise Exception(f"{method} {url} {response.status} {response.read()}")
              def offboard(self):
                  try:
                      self.delete_account()
                  except:
                      print("deleting account from DX failed, but passing the event")
          def enroll_for_compute_optimizer():
              print("enrolling for compute optimizer")
              try:
                  client = boto3.client("compute-optimizer")
                  response = client.get_enrollment_status()
                  print("current enrollment status : " + str(response))
                  if response["status"] == "Active":
                      return print("already enrolled")
                  response = client.update_enrollment_status(status="Active")
                  print("updated enrollment status : " + str(response))
              except Exception as err:
                  print("error occurred while enrolling compute-optimizer" + str(err))
          def handler(event, context):
              req_type = event.get("RequestType")
              print("request type: " + req_type)
              if req_type not in ["Create", "Update"]:
                  print("No action required for requests other than `Create` or `Update`")
                  return cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              try:
                  params = event["ResourceProperties"]
                  print(f"payload : {json.dumps(params)}")
                  if params["Operation"] == "SaveAccountDetails":
                      Onboarding(params).onboard()
                  elif params["Operation"] == "EnrollForComputeOptimizer":
                      enroll_for_compute_optimizer()
                  return cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as err:
                  print(traceback.format_exc())
                  return cfnresponse.send(event, context, cfnresponse.FAILED, {}, reason=f"err={err}, cloudwatch-log-stream={context.log_stream_name}")
      Tags: [{ Key: "app", Value: "Mavvrik" }, { Key: "version", Value: "v3" }, { Key: "variant", Value: "v3" }]

  SaveAccountDetails:
    Type: AWS::CloudFormation::CustomResource
    Condition: TokenAvailable
    DependsOn: EnrollForComputeOptimizer
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
      Operation: "SaveAccountDetails"
      Host: "api.mavvrik.ai"
      Token: !Ref token
      AccessType: !Ref accessType
      Tenant: !Ref tenant
      StackId: !If [ParentAvailable, !Ref parentStackId, !Ref "AWS::StackId"]
      RoleARN: !GetAtt ResourceDiscoveryRole.Arn
      AccountId: !Ref "AWS::AccountId"
      OrgUnitIds: !Ref parentStackOrgUnitIds

  EnrollForComputeOptimizer:
    Type: AWS::CloudFormation::CustomResource
    Condition: TokenAvailable
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
      Operation: "EnrollForComputeOptimizer"

Outputs:
  AccountId:
    Value: !Ref "AWS::AccountId"
  StackId:
    Value: !If [ParentAvailable, !Ref parentStackId, !Ref "AWS::StackId"]
  RoleARN:
    Value: !GetAtt ResourceDiscoveryRole.Arn