AWSTemplateFormatVersion: '2010-09-09'
Description: 'TrailTool Producer (Logs Account) - S3 bucket configuration and EventBridge forwarding'

Parameters:
  CloudTrailBucketName:
    Type: String
    Description: Name of the S3 bucket containing CloudTrail logs

  TrailToolAccountId:
    Type: String
    Description: AWS Account ID where TrailTool is deployed (SaaS account)
    Default: '468087121425'

  UserEmail:
    Type: String
    Description: Email address for TrailTool registration and notifications


Resources:
  # ===================================================================
  # CUSTOM RESOURCE LAMBDA - Adds bucket policy and enables EventBridge
  # ===================================================================
  
  BucketPolicyHelperRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: BucketPolicyAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - s3:GetBucketPolicy
                  - s3:PutBucketPolicy
                  - s3:PutBucketNotificationConfiguration
                  - s3:GetBucketNotificationConfiguration
                  - s3:GetBucketNotification
                  - s3:PutBucketNotification
                Resource: !Sub 'arn:aws:s3:::${CloudTrailBucketName}'

  BucketPolicyHelperFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: trailtool-bucket-config
      Runtime: python3.12
      Handler: index.handler
      Role: !GetAtt BucketPolicyHelperRole.Arn
      Timeout: 60
      Code:
        ZipFile: |
          import json
          import boto3
          from urllib.request import Request, urlopen

          s3 = boto3.client('s3')

          # Inline cfnresponse since it's not available in Python 3.12 runtime
          SUCCESS = "SUCCESS"
          FAILED = "FAILED"

          def send(event, context, responseStatus, responseData, physicalResourceId=None, noEcho=False, reason=None):
              responseUrl = event['ResponseURL']
              
              responseBody = {
                  'Status': responseStatus,
                  'Reason': reason or "See the details in CloudWatch Log Stream: {}".format(context.log_stream_name),
                  'PhysicalResourceId': physicalResourceId or context.log_stream_name,
                  'StackId': event['StackId'],
                  'RequestId': event['RequestId'],
                  'LogicalResourceId': event['LogicalResourceId'],
                  'NoEcho': noEcho,
                  'Data': responseData
              }
              
              json_responseBody = json.dumps(responseBody)
              
              headers = {
                  'content-type': '',
                  'content-length': str(len(json_responseBody))
              }
              
              try:
                  request = Request(responseUrl, data=json_responseBody.encode('utf-8'), headers=headers, method='PUT')
                  response = urlopen(request)
                  print("Status code:", response.getcode())
                  print("Status message:", response.msg)
              except Exception as e:
                  print("send(..) failed executing request:", e)

          def handler(event, context):
              try:
                  print(f"Event: {json.dumps(event)}")
                  
                  bucket = event['ResourceProperties']['BucketName']
                  account_id = event['ResourceProperties']['TrailToolAccountId']
                  
                  if event['RequestType'] == 'Delete':
                      # On delete, remove our statement from the policy
                      try:
                          policy = s3.get_bucket_policy(Bucket=bucket)
                          policy_doc = json.loads(policy['Policy'])
                          
                          # Remove TrailToolReadAccess statement
                          policy_doc['Statement'] = [
                              stmt for stmt in policy_doc['Statement'] 
                              if stmt.get('Sid') != 'TrailToolReadAccess'
                          ]
                          
                          s3.put_bucket_policy(
                              Bucket=bucket,
                              Policy=json.dumps(policy_doc)
                          )
                          print(f"Removed TrailToolReadAccess from bucket {bucket}")
                      except Exception as e:
                          print(f"Error removing policy: {e}")
                          # Don't fail on delete
                      
                      send(event, context, SUCCESS, {})
                      return
                  
                  # Create or Update: Configure bucket
                  if event['RequestType'] in ['Create', 'Update']:
                      # Enable EventBridge on S3 bucket (idempotent, preserves existing notifications)
                      try:
                          current_config = s3.get_bucket_notification_configuration(Bucket=bucket)
                          # Remove ResponseMetadata if present
                          current_config.pop('ResponseMetadata', None)
                          
                          # Add EventBridge if not already present
                          if 'EventBridgeConfiguration' not in current_config:
                              current_config['EventBridgeConfiguration'] = {}
                              s3.put_bucket_notification_configuration(
                                  Bucket=bucket,
                                  NotificationConfiguration=current_config
                              )
                              print(f"Enabled EventBridge on bucket {bucket}")
                          else:
                              print(f"EventBridge already enabled on bucket {bucket}")
                      except Exception as e:
                          print(f"Error enabling EventBridge: {e}")
                          # Don't fail the whole stack if EventBridge enablement fails
                      
                      # Add bucket policy statement
                      try:
                          policy = s3.get_bucket_policy(Bucket=bucket)
                          policy_doc = json.loads(policy['Policy'])
                      except s3.exceptions.NoSuchBucketPolicy:
                          # No policy exists, create one
                          policy_doc = {
                              "Version": "2012-10-17",
                              "Statement": []
                          }
                      
                      # Remove existing TrailToolReadAccess if present
                      policy_doc['Statement'] = [
                          stmt for stmt in policy_doc['Statement'] 
                          if stmt.get('Sid') != 'TrailToolReadAccess'
                      ]
                      
                      # Add our statement
                      trailtool_statement = {
                          "Sid": "TrailToolReadAccess",
                          "Effect": "Allow",
                          "Principal": {
                              "AWS": f"arn:aws:iam::{account_id}:root"
                          },
                          "Action": [
                              "s3:GetObject",
                              "s3:ListBucket"
                          ],
                          "Resource": [
                              f"arn:aws:s3:::{bucket}",
                              f"arn:aws:s3:::{bucket}/*"
                          ]
                      }
                      
                      policy_doc['Statement'].append(trailtool_statement)
                      
                      # Put the updated policy
                      s3.put_bucket_policy(
                          Bucket=bucket,
                          Policy=json.dumps(policy_doc)
                      )
                      
                      print(f"Successfully added TrailToolReadAccess to bucket {bucket}")
                      send(event, context, SUCCESS, {})
                  
              except Exception as e:
                  print(f"Error: {e}")
                  send(event, context, FAILED, {"Error": str(e)})

  # ===================================================================
  # S3 BUCKET POLICY - Grant TrailTool read access
  # ===================================================================

  TrailToolBucketConfiguration:
    Type: Custom::BucketConfiguration
    Properties:
      ServiceToken: !GetAtt BucketPolicyHelperFunction.Arn
      BucketName: !Ref CloudTrailBucketName
      TrailToolAccountId: !Ref TrailToolAccountId

  # ===================================================================
  # EVENTBRIDGE RULE - Invoke TrailTool ingestor on S3 events
  # ===================================================================
  # Note: EventBridge is enabled on the bucket by the BucketPolicyHelperFunction

  CloudTrailEventsRule:
    Type: AWS::Events::Rule
    DependsOn: TrailToolBucketConfiguration
    Properties:
      Name: trailtool-cloudtrail-forwarding
      Description: Invoke TrailTool ingestor Lambda on CloudTrail S3 events
      EventPattern:
        source:
          - aws.s3
        detail-type:
          - Object Created
        detail:
          bucket:
            name:
              - !Ref CloudTrailBucketName
          object:
            key:
              - prefix: AWSLogs/
              - suffix: .json.gz
      State: ENABLED
      Targets:
        - Arn: !Sub 'arn:aws:lambda:${AWS::Region}:${TrailToolAccountId}:function:trailtool-ingestor'
          Id: TrailToolIngestor
          RoleArn: !GetAtt EventBridgeInvokeRole.Arn

  # IAM Role for EventBridge to invoke cross-account Lambda
  EventBridgeInvokeRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: TrailToolEventBridgeInvokeRole
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: events.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: InvokeTrailToolLambda
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: lambda:InvokeFunction
                Resource: !Sub 'arn:aws:lambda:${AWS::Region}:${TrailToolAccountId}:function:trailtool-ingestor'

  # ===================================================================
  # REGISTRATION WITH TRAILTOOL
  # ===================================================================

  TrailToolRegistration:
    Type: Custom::TrailToolRegistration
    DependsOn:
      - CloudTrailEventsRule
    Properties:
      ServiceToken: !Sub 'arn:aws:lambda:${AWS::Region}:${TrailToolAccountId}:function:trailtool-onboard'
      UserEmail: !Ref UserEmail
      CloudTrailBucket: !Ref CloudTrailBucketName
      CloudTrailAccount: !Ref 'AWS::AccountId'
      ManagementAccountId: !Ref 'AWS::AccountId'
      EventBridgeRuleArn: !GetAtt CloudTrailEventsRule.Arn

Outputs:
  CustomerId:
    Description: Your TrailTool customer ID
    Value: !GetAtt TrailToolRegistration.CustomerId
  
  CloudTrailBucketName:
    Description: CloudTrail bucket with EventBridge enabled
    Value: !Ref CloudTrailBucketName
  
  EventBridgeRuleArn:
    Description: ARN of the EventBridge rule forwarding to TrailTool
    Value: !GetAtt CloudTrailEventsRule.Arn
