Resources:
  # creates a glue database for flow logs athena tables.
  VpcFlowLogsAthenaDatabase:
    Type: AWS::Glue::Database
    Properties:
      DatabaseInput:
        Name: 'vpcflowlogsathenadatabasefl0dd90b1f0faeadb4d'
      CatalogId: !Ref AWS::AccountId

  # this section creates glue athenaTables mapped to s3 flow logs bucket
    # creates a workgroup in Athena for queries.
  fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup:
    Type: AWS::Athena::WorkGroup
    Properties:
      Name: fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
      Description: This workgroup has the queries related to vpc flow logs.
      State: ENABLED
      Tags:
        - Key: "service"
          Value: "vpc_flow_logs"
      WorkGroupConfiguration:
        BytesScannedCutoffPerQuery: 200000000
        EnforceWorkGroupConfiguration: true
        PublishCloudWatchMetricsEnabled: true
        RequesterPaysEnabled: true
        ResultConfiguration:
          OutputLocation: s3://lahora-athena-queries
        EngineVersion:
          SelectedEngineVersion: 'Athena engine version 2'

  VpcFlowLogsTotalBytesfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Top 50 pairs of source and destination IP addresses with the most bytes transferred. The start and end in the WHERE clause are placeholders"
      Name: VpcFlowLogsTotalBytesTransferredfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT SUM(bytes) as totalbytes, srcaddr, dstaddr from fl0dd90b1f0faeadb4ddaily2022030920220309
        WHERE "start"=1646832775562 AND "end"=1646832775562
        GROUP BY srcaddr, dstaddr
        ORDER BY totalbytes
        LIMIT 50
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup


  VpcFlowLogsTopTalkersfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "The top 50 IPs transmitting traffic in your VPC/ENI/Subnet."
      Name: VpcFlowLogsTopTalkersfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT srcaddr, sum(bytes) as bytes
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        group by srcaddr
        order by bytes desc
        limit 50
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup


  VpcFlowLogsSshRdpTrafficfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Monitor SSH and RDP traffic"
      Name: VpcFlowLogsSshRdpTrafficfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT *
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        WHERE srcport in (22,3389) OR dstport IN (22, 3389)
        ORDER BY "start" ASC
        limit 50
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup

  VpcFlowLogsAdminPortTrafficfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Monitor the traffic on administrative web app ports"
      Name: VpcFlowLogsAdminPortTrafficfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT ip, sum(bytes) as total_bytes
        FROM (
        SELECT dstaddr as ip,sum(bytes) as bytes
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        GROUP BY 1
        UNION ALL
        SELECT srcaddr as ip,sum(bytes) as bytes
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        GROUP BY 1
        )
        GROUP BY ip
        ORDER BY total_bytes
        DESC LIMIT 10
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup

  VpcFlowLogsRejectedTrafficfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Recorded traffic which was not permitted by the security groups or network ACLs."
      Name: VpcFlowLogsRejectedTrafficfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT srcaddr, dstaddr,  count(*) count, "action"
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        WHERE "action" = 'REJECT'
        GROUP BY srcaddr, dstaddr, "action"
        ORDER BY count desc
        LIMIT 25
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup



  VpcFlowLogsTrafficFrmSrcAddrfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Traffic transmitted from a particular source IP address. Here the source address (srcaddr) is a placeholder"
      Name: VpcFlowLogsTrafficFrmSrcAddrfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT *
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        WHERE srcaddr = '198.51.100.2'
        ORDER BY "start" ASC
        LIMIT 50;
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup

  VpcFlowLogsTrafficToDstAddrfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Athena::NamedQuery
    DependsOn:
      - VpcFlowLogsAthenaDatabase
      - fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup
    Properties:
      Database: !Ref VpcFlowLogsAthenaDatabase
      Description: "Traffic transmitted from a particular destination IP address. Here the destination address (dstaddr) is a placeholder"
      Name: VpcFlowLogsTrafficToDstAddrfl0dd90b1f0faeadb4ddaily2022030920220309
      QueryString: >
        SELECT *
        FROM fl0dd90b1f0faeadb4ddaily2022030920220309
        WHERE dstaddr = '198.51.100.2'
        ORDER BY "start" ASC
        LIMIT 50;
      WorkGroup: !Ref fl0dd90b1f0faeadb4ddaily2022030920220309WorkGroup






  VpcFlowLogsAthenafl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Glue::Table
    DependsOn:
      - VpcFlowLogsAthenaDatabase
    Properties:
      CatalogId: !Ref AWS::AccountId
      DatabaseName: !Ref VpcFlowLogsAthenaDatabase
      TableInput:
        Description: This table has the schema for vpc flow logs information.
        Name: fl0dd90b1f0faeadb4ddaily2022030920220309
        PartitionKeys:
          - Name: year
            Type: string
          - Name: month
            Type: string
          - Name: day
            Type: string
        TableType: EXTERNAL_TABLE
        StorageDescriptor:
          Location:  s3://lahora-vpc-flow-logs/AWSLogs/801915948376/vpcflowlogs/us-east-1/
          InputFormat: org.apache.hadoop.mapred.TextInputFormat
          OutputFormat: org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat
          SerdeInfo:
            Parameters:
              skip.header.line.count: "1"
              EXTERNAL: "true"
              field.delim: ' '
              serialization.format: ' '
            SerializationLibrary: org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe
          Columns:
            - Name: 'version'
              Type: int
            - Name: 'account_id'
              Type: string
            - Name: 'interface_id'
              Type: string
            - Name: 'srcaddr'
              Type: string
            - Name: 'dstaddr'
              Type: string
            - Name: 'srcport'
              Type: int
            - Name: 'dstport'
              Type: int
            - Name: 'protocol'
              Type: bigint
            - Name: 'packets'
              Type: bigint
            - Name: 'bytes'
              Type: bigint
            - Name: 'start'
              Type: bigint
            - Name: 'end'
              Type: bigint
            - Name: 'action'
              Type: string
            - Name: 'log_status'
              Type: string

  # creates an IAM role for lambda function execution permission.
  VPCFlowLogsAthenaIntegrationLambdaExecutorRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: VPCFlowLogsAthenaIntegrationLambdaExecutorPolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
              - Effect: Allow
                Action:
                  - 'glue:GetTable'
                  - 'glue:CreatePartition'
                  - 'glue:UpdatePartition'
                  - 'glue:GetPartition'
                Resource: "*"

  # creates a lambda function to initialize the partitions based on the config start, end date and frequency.
  InitializerAsync:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          const response = require('./cfn-response');
          const AWS = require('aws-sdk');
          const glue = new AWS.Glue();

          exports.handler =  async function(event, context) {
            let errs = [], status
            if (event.RequestType === 'Delete') {
              status = response.SUCCESS
            } else {
                console.log("Parsing athena configs")
                let rp = event.ResourceProperties
                let confs = rp.athenaIntegrations;
                let db = rp.dbName
                let hive = rp.hive
                let account_id = rp.account_id
                let service = rp.service
                let region = rp.region

                let errs = []

                for(let i = 0; i < confs.length; i++) {
                  let cnf = confs[i]
                  let tab = cnf['partitionTableName']
                  let frq = cnf['partitionLoadFrequency']
                  let strt = (cnf['partitionStartDate'] == undefined) ? new Date() : new Date(cnf['partitionStartDate'])
                  let end = (cnf['partitionEndDate'] == undefined) ? new Date() : new Date(cnf['partitionEndDate'])

                  while(strt <= end) {
                    let table = await glue.getTable({
                      DatabaseName: db,
                      Name: tab,
                    }).promise()

                    let strgDesc = table.Table['StorageDescriptor']
                    let Values
                    let newDate = new Date()

                    if(frq == "monthly") {
                      if(hive == "true") {
                        Values = ["aws-account-id=" + account_id, "aws-service=" + service, "aws-region=" + region, "year=" + String(strt.getFullYear()), "month=" + ("0" + (strt.getMonth() + 1)).slice(-2)]
                      } else {
                        Values = [String(strt.getFullYear()), ("0" + (strt.getMonth() + 1)).slice(-2)]
                      }
                      newDate = strt.setMonth(strt.getMonth() + 1);
                    } else if(frq == "hourly") {
                      if(hive == "true") {
                        Values = ["aws-account-id=" + account_id, "aws-service=" + service, "aws-region=" + region, "year=" + String(strt.getFullYear()), "month=" + ("0" + (strt.getMonth() + 1)).slice(-2), "day=" + ("0" + strt.getDate()).slice(-2), "hour=" + ("0" + strt.getHours()).slice(-2)]
                      } else {
                        Values = [String(strt.getFullYear()), ("0" + (strt.getMonth() + 1)).slice(-2), ("0" + strt.getDate()).slice(-2), ("0" + strt.getHours()).slice(-2)]
                      }
                      newDate.setHours(strt.getHours() + 1);
                    } else {
                      if(hive == "true") {
                        Values = ["aws-account-id=" + account_id, "aws-service=" + service, "aws-region=" + region, "year=" + String(strt.getFullYear()), "month=" + ("0" + (strt.getMonth() + 1)).slice(-2), "day=" + ("0" + strt.getDate()).slice(-2)]
                      } else {
                        Values = [String(strt.getFullYear()), ("0" + (strt.getMonth() + 1)).slice(-2), ("0" + strt.getDate()).slice(-2)]
                      }
                      newDate = strt.setDate(strt.getDate() + 1);
                    }

                    try {
                      let result = await glue.getPartition({
                          DatabaseName: db,
                          TableName: tab,
                          PartitionValues: Values
                      }).promise()
                    } catch (err) {
                      if(err.code === 'EntityNotFoundException'){
                          console.log(strgDesc)
                          let params = {
                              DatabaseName: db,
                              TableName: tab,
                              PartitionInput: {
                                  StorageDescriptor: {
                                      ...strgDesc,
                                      Location: `${strgDesc.Location}${Values.join('/')}/`
                                  },
                                  Values,
                              },
                          }
                          await glue.createPartition(params).promise()
                      } else {
                          errs.push(err)
                      }
                    }
                    strt = new Date(newDate);
                  }
                }

                status = errs.length > 0 ? response.FAILED : response.SUCCESS
              }
              return new Promise(() => response.send(event, context, status,
              errs.length > 0 ? { error: errs } : {}, event.LogicalResourceId));
          }

      Handler: 'index.handler'
      Timeout: 60
      Runtime: nodejs14.x
      ReservedConcurrentExecutions: 1
      Role: !GetAtt VPCFlowLogsAthenaIntegrationLambdaExecutorRole.Arn

  # creates an initial custom trigger for lambda function for partition creation.
  Initializer2022030920220309:
    Type: 'Custom::VPCFlowLogsAthenaStartInitializer'
    DependsOn:
      - InitializerAsync
      - VpcFlowLogsAthenafl0dd90b1f0faeadb4ddaily2022030920220309
    Properties:
      ServiceToken: !GetAtt InitializerAsync.Arn
      dbName: !Ref VpcFlowLogsAthenaDatabase
      hive : false
      account_id : "801915948376"
      service: vpcflowlogs
      region: us-east-1
      athenaIntegrations:
        - partitionTableName: fl0dd90b1f0faeadb4ddaily2022030920220309
          partitionLoadFrequency: daily
          partitionStartDate: 2022-03-09T13:32:55
          partitionEndDate: 2022-03-09T13:32:55

  # creates a lambda function for daily partition creation.
  Partitionerfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Lambda::Function
    DependsOn:
      - VpcFlowLogsAthenafl0dd90b1f0faeadb4ddaily2022030920220309
    Properties:
      Code:
        ZipFile: |
          const AWS = require('aws-sdk');
          const glue = new AWS.Glue();

          exports.handler = async function(event, context) {
            let db = event.db
            let confs = event.athena
            let hive = event.hive
            let account_id = event.account_id
            let service = event.service
            let region = event.region
            let today = new Date()
            let errs = []

            for(let i = 0; i < confs.length; i++) {
              let cnf = confs[i]
              let tab = cnf['partitionTableName']
              let frq = cnf['frequency']

              let table = await glue.getTable({
                DatabaseName: db,
                Name: tab,
              }).promise()

              let strgDesc = table.Table['StorageDescriptor']
              let Values

              if(frq == "hourly"){
                if(hive == "true") {
                  Values = ["aws-account-id=" + account_id, "aws-service=" + service, "aws-region=" + region, "year=" + String(today.getFullYear()), "month=" + ("0" + (today.getMonth() + 1)).slice(-2), "day=" + ("0" + today.getDate()).slice(-2), "hour=" + ("0" + today.getHours()).slice(-2)]
                } else {
                  Values = [String(today.getFullYear()), ("0" + (today.getMonth() + 1)).slice(-2), ("0" + today.getDate()).slice(-2), ("0" + today.getHours()).slice(-2)]
                }
              } else {
                  if(hive == "true") {
                    Values = ["aws-account-id=" + account_id, "aws-service=" + service, "aws-region=" + region, "year=" + String(today.getFullYear()), "month=" + ("0" + (today.getMonth() + 1)).slice(-2), "day=" + ("0" + (today.getDate())).slice(-2)]
                  } else {
                    Values = [String(today.getFullYear()), ("0" + (today.getMonth() + 1)).slice(-2), ("0" + (today.getDate())).slice(-2)]
                  }
              }
              try {
                let result = await glue.getPartition({
                  DatabaseName: db,
                  TableName: tab,
                  PartitionValues: Values
                }).promise()
              } catch (err) {
                  if(err.code === 'EntityNotFoundException'){
                    console.log(strgDesc)
                    let params = {
                      DatabaseName: db,
                      TableName: tab,
                      PartitionInput: {
                        StorageDescriptor: {
                            ...strgDesc,
                            Location: `${strgDesc.Location}${Values.join('/')}/`
                        },
                        Values,
                      },
                    }
                    await glue.createPartition(params).promise()
                  } else {
                    errs.push(err)
                }
              }
            }

            return new Promise(function(resolve, reject) { errs.length > 0 ? reject(errs) : resolve("SUCCESS")});
          }
      Handler: 'index.handler'
      Timeout: 30
      Runtime: nodejs14.x
      ReservedConcurrentExecutions: 1
      Role: !GetAtt VPCFlowLogsAthenaIntegrationLambdaExecutorRole.Arn

  # creates event rule for daily lambda function trigger
  ScheduledEventRulefl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Events::Rule
    DependsOn: Partitionerfl0dd90b1f0faeadb4ddaily2022030920220309
    Properties:
      Description: This event rule will be invoking lambda based on schedule
      Name: ScheduledEventRulefl0dd90b1f0faeadb4ddaily2022030920220309
      ScheduleExpression: cron(0 0 * * ? *)
      State: ENABLED
      Targets:
        - Arn: !GetAtt Partitionerfl0dd90b1f0faeadb4ddaily2022030920220309.Arn
          Id: Partitionerfl0dd90b1f0faeadb4ddaily2022030920220309
          Input: "{\"db\": \"vpcflowlogsathenadatabasefl0dd90b1f0faeadb4d\", \"hive\": \"false\", \"account_id\": \"801915948376\", \"service\": \"vpcflowlogs\", \"region\": \"us-east-1\" ,\"athena\": [ {\"partitionTableName\": \"fl0dd90b1f0faeadb4ddaily2022030920220309\", \"frequency\": \"daily\"}]}"

  # creates lambda permission for daily schedule
  ScheduledEventPermissionfl0dd90b1f0faeadb4ddaily2022030920220309:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !GetAtt Partitionerfl0dd90b1f0faeadb4ddaily2022030920220309.Arn
      Principal: events.amazonaws.com
      SourceArn: !GetAtt ScheduledEventRulefl0dd90b1f0faeadb4ddaily2022030920220309.Arn

