#!/bin/python
import boto3
from botocore.errorfactory import ClientError
from pprint import pprint, pformat
import json
import sys
import argparse
import logging


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--profile', help='AWS profile', default='gdb-delphi-dev')
    parser.add_argument('--region', help='AWS Default region', default='us-east-1')
    parser.add_argument('--keyalias', help='KMS key for GD export', default='alias/guardduty')
    return parser.parse_args()

def find_key_arn_by_alias(client,alias):
    aliases = client.list_aliases()['Aliases']
    key_id = list(filter(lambda x: x['AliasName'] == alias, aliases))[0]['TargetKeyId']
    return client.describe_key(KeyId=key_id)['KeyMetadata']['Arn']

def create_dir_if_not_exist(client,bucket,path):
    try:
        client.head_object(Bucket=bucket, Key=path)
    except ClientError:
        client.put_object(Bucket=bucket,Body='',Key=path)

if __name__ == "__main__":
    logger = logging.getLogger('')
    logger.setLevel(logging.DEBUG)

    args = parse_arguments()
    regions = [
        "us-east-1",
        "us-east-2",
        "us-east-1",
        "us-west-1",
        "us-west-2",
        "ap-south-1",
        "ap-northeast-2",
        "ap-southeast-1",
        "ap-southeast-2",
        "ap-northeast-1",
        "ca-central-1",
        "eu-central-1",
        "eu-west-1",
        "eu-west-2",
        "eu-west-3",
        "eu-north-1",
        "sa-east-1",
    ]

    # Create GuardDuty clients for all regions and s3 and kms client for default region
    if args.profile:
        gd_clients = { x : boto3.Session(profile_name=args.profile,region_name=x).client('guardduty') for x in regions }
        account_id = boto3.Session(profile_name=args.profile,region_name=args.region).client('sts').get_caller_identity()["Account"]
        kms_client = boto3.Session(profile_name=args.profile,region_name=args.region).client('kms')
        s3_client = boto3.Session(profile_name=args.profile,region_name=args.region).client('s3')
    else:
        gd_clients = { x : boto3.client('guardduty',profile_name=args.profile,region_name=x) for x in regions }
        account_id = boto3.client('sts',profile_name=args.profile,region_name=args.region).get_caller_identity()["Account"]
        kms_client = boto3.client('kms',profile_name=args.profile,region_name=args.region)
        s3_client = boto3.client('s3',profile_name=args.profile,region_name=args.region)

    export_bucket = f"arn:aws:s3:::security-{account_id}-logs/gd/"
    key_arn = find_key_arn_by_alias(kms_client, args.keyalias)
    create_dir_if_not_exist(s3_client,f"security-{account_id}-logs","gd/")

    for region in regions:
        destinations = []
        detectorIds = gd_clients[region].list_detectors()['DetectorIds']
        for detectorId in detectorIds:
            if (len(gd_clients[region].list_publishing_destinations(DetectorId=detectorId)['Destinations']) == 0):
                gd_clients[region].create_publishing_destination(
                    DetectorId=detectorId,
                    DestinationType='S3',
                    DestinationProperties={
                        'DestinationArn': export_bucket,
                        'KmsKeyArn': key_arn
                    }
                )
            else:
                print("destination already exists in this region")
            destinations += gd_clients[region].list_publishing_destinations(DetectorId=detectorId)['Destinations']
        pprint(destinations)