"""Athena Query
Submits the appropriate Security Fairy
query to Athena.
"""


import re
import boto3
import logging
from datetime import datetime, timedelta
from setup_logger import create_logger
from botocore.exceptions import ProfileNotFound

logger = create_logger(name="athena_query.py", logging_level=logging.INFO)

try:
    SESSION = boto3.session.Session(profile_name='training',
                                    region_name='us-east-1')
except ProfileNotFound as pnf:
    SESSION = boto3.session.Session()


def lambda_handler(event, context):
    """ Executed by the Lambda service.
    Submits the query for execution and returns
    the Execution ID for use by subsequent
    Lambda functions.
    """

    event['execution_id'] = execute_query(event['entity_arn'],
                                          event['num_days'],
                                          event['s3_bucket'])

    return event


def window_calc(num_days):
    """Calculate the correct year,
    month, and day for the query
    """
    delta = timedelta(days=num_days)
    today = datetime.now()
    query_date = today - delta
    year = query_date.year
    month = query_date.month

    return year, month


def execute_query(entity_arn, num_days, s3_bucket):
    """Submit and run query"""

    escaped_arn = build_escaped_arn(entity_arn)
    year, month = window_calc(num_days)

    hql = f"""
       select useridentity.arn as user_arn
            , eventsource
            , array_distinct(array_agg(eventName)) as actions
         from aws_logs.cloudtrail
        where year = '{year}'
          and month >= '{month}'
          and regexp_like(useridentity.arn, '{escaped_arn}\/.+')
     group by useridentity.arn
            , eventsource
          """
    logger.info(hql)

    output = f's3://{s3_bucket}/tables'
    config = {
        'OutputLocation': output,
        'EncryptionConfiguration': {
            'EncryptionOption': 'SSE_S3'
        }
    }

    athena_client = SESSION.client('athena')
    execution = athena_client.start_query_execution(QueryString=hql,
                                                    ResultConfiguration=config)
    logger.info("Query ID:")
    logger.info(execution['QueryExecutionId'])

    return execution['QueryExecutionId']


def build_escaped_arn(entity_arn):
    """Format ARN"""

    split_arn = re.split('/|:', entity_arn)
    escaped_arn = "arn:aws:sts::" + split_arn[4] + ":assumed-role\\/" + split_arn[6]
    logger.debug(escaped_arn)
    return escaped_arn

if __name__ == '__main__':
    # arn:aws:sts::281782457076:assumed-role\/1s_tear_down_role\/.+
    lambda_handler(
        {
            "entity_arn": "arn:aws:iam::281782457076:assumed-role/1s_tear_down_role",
            "num_days": "-30",
            "s3_bucket": "1s-potato-east"
        },
        {}
    )

@@ -1015,24 +1015,49 @@
 ery%0A    %22%22%22%0A
+    days = abs(num_days)%0A
     delta = 
@@ -1071,20 +1071,16 @@
 ta(days=
-num_
 days)%0A  
@@ -2627,16 +2627,38 @@
 %5C/.+%0A   
+ execute_query()%0A    #
  lambda_
@@ -2665,24 +2665,26 @@
 handler(%0A   
+ #
      %7B%0A     
@@ -2677,24 +2677,26 @@
  #     %7B%0A   
+ #
          %22en
@@ -2770,24 +2770,25 @@
 e%22,%0A    
+#
         
 %22num_day
@@ -2779,16 +2779,17 @@
         
+ 
 %22num_day
@@ -2798,24 +2798,26 @@
 : %22-30%22,%0A   
+ #
          %22s3
@@ -2842,24 +2842,26 @@
 to-east%22%0A   
+ #
      %7D,%0A    
@@ -2855,24 +2855,26 @@
 #     %7D,%0A   
+ #
      %7B%7D%0A    
@@ -2872,11 +2872,13 @@
   %7B%7D%0A   
+ #
  )%0A

change window_calc to use absolute value of num_days