#!/usr/bin/env python
# -*- coding: utf-8  -*-

import boto3
from botocore.exceptions import ClientError
import botocore
import logging
import csv
from argparse import ArgumentParser
from collections import OrderedDict

logger = logging.getLogger()
logging.basicConfig()
logging.getLogger('boto3').setLevel(logging.ERROR)
logging.getLogger('boto3').propagate = False
logging.getLogger('botocore').setLevel(logging.ERROR)
logging.getLogger('botocore').propagate = False

parser = ArgumentParser("export-import-user-pool-users")

parser.add_argument(
    "-r",
    "--region",
    dest="regionName",
    help="Cognito Region (required)",
    required=True)
parser.add_argument(
    "-c",
    "--cognito-pool-id",
    dest="userPoolId",
    help="Cognito Pool Id (required)",
    required=True)
parser.add_argument(
    "-e",
    "--exports",
    action="store_true",
    help="Switch to enable export users, groups, users-groups map to files",
    required=False)
parser.add_argument(
    "-i",
    "--imports",
    action="store_true",
    help="Switch to enable import users, groups, users-groups map to files",
    required=False)
parser.add_argument(
    "-f",
    "--file",
    dest="usersFile",
    help="CSV file with users data",
    required=False)
parser.add_argument(
    "-g",
    "--groups-file",
    dest="groupsFile",
    help="CSV file with groups data",
    required=False)
parser.add_argument(
    "-m",
    "--map-file",
    dest="usersGroupsFile",
    help="CSV file with users-groups map",
    required=False)
parser.add_argument(
    "-l",
    "--logger",
    dest="logger",
    help="Log level",
    choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
    default='INFO',
    required=False)

pargs = parser.parse_args()

logger.setLevel(pargs.logger)


class AwsCognito(object):
    def __init__(self):
        super(AwsCognito, self).__init__()
        self.boto3 = boto3.client('cognito-idp', region_name=pargs.regionName)
        self.csvHeaders = self.getHeaderCSV()
        self.userTemplate = self.userTemplate()

    def _params(self, paginationToken: str = '', groupName: str = ''):
        params = {'UserPoolId': pargs.userPoolId}
        if paginationToken != '':
            params['PaginationToken'] = paginationToken
        if groupName != '':
            params['GroupName'] = groupName
        return params

    def getUsers(self, paginationToken: str = ''):
        return self.boto3.list_users(**self._params(paginationToken))

    def getGroups(self, paginationToken: str = ''):
        return self.boto3.list_groups(**self._params(paginationToken))

    def getUsersWithinGroup(self,
                            paginationToken: str = '',
                            groupName: str = ''):
        params = self._params(paginationToken=paginationToken,
                              groupName=groupName)
        if 'PaginationToken' in params:
            params['NextToken'] = params.pop('PaginationToken')
        return self.boto3.list_users_in_group(
            **params)

    def getHeaderCSV(self):
        try:
            return self.boto3.get_csv_header(
                UserPoolId=pargs.userPoolId)
        except Exception as e:
            logger.error(e)
            exit()

    def userTemplate(self):
        return OrderedDict({attr: ''
                            for attr in self.csvHeaders['CSVHeader']})

    def _mergeUserTemplateDicts(self, user):
        userAttrs = user['Attributes']
        mergedUserTemplate = self.userTemplate.copy()

        for key in userAttrs:
            if key['Name'] in self.csvHeaders['CSVHeader']:
                mergedUserTemplate[key['Name']] = key['Value']

        mergedUserTemplate['cognito:username'] = user['Username']

        if 'MFAOptions' in user:
            mergedUserTemplate['cognito:mfa_enabled'] = 'true'
        else:
            mergedUserTemplate['cognito:mfa_enabled'] = 'false'

        if mergedUserTemplate['phone_number_verified'] == '':
            mergedUserTemplate['phone_number_verified'] = 'false'

        if mergedUserTemplate['email_verified'] == '':
            mergedUserTemplate['email_verified'] = 'false'

        return mergedUserTemplate

    def _traversUsers(self):
        params = {}

        while True:
            usersData = self.getUsers(**params)

            if 'Users' in usersData:
                for user in usersData['Users']:
                    yield self._mergeUserTemplateDicts(user)

            if 'PaginationToken' in usersData:
                params['PaginationToken'] = usersData['PaginationToken']
            else:
                break

    def storeUsersToFile(self):
        if pargs.usersFile:
            fileName = pargs.usersFile
        else:
            fileName = pargs.userPoolId + '-users.csv'

        with open(fileName, 'w', newline='') as csvFile:
            self.csv = csv.DictWriter(csvFile,
                                      fieldnames=self.userTemplate,
                                      delimiter=',')

            self.csv.writeheader()

            for user in self._traversUsers():
                self.csv.writerow(user)

            csvFile.close()

    def _traverseGroups(self):
        params = {}

        while True:
            groupsData = self.getGroups(**params)

            if 'Groups' in groupsData:
                for group in groupsData['Groups']:
                    yield group['GroupName']

            if 'PaginationToken' in groupsData:
                params['PaginationToken'] = groupsData['PaginationToken']
            else:
                break

    def _traverseUsersWithinGroup(self, groupName: str = ''):
        params = {}
        params['groupName'] = groupName

        while True:
            usersData = self.getUsersWithinGroup(**params)

            if 'Users' in usersData:
                for user in usersData['Users']:
                    yield {'groupName': groupName,
                           'userName': user['Username']}

            if 'NextToken' in usersData:
                params['PaginationToken'] = usersData['PaginationToken']
            else:
                break

    def storeUsersGroupsMapToFile(self):
        fieldnames = ['groupName', 'userName']
        if pargs.usersFile:
            fileName = pargs.usersGroupsFile
        else:
            fileName = pargs.userPoolId + '-map.csv'

        with open(fileName, 'w', newline='') as csvFile:
            self.csv = csv.DictWriter(csvFile,
                                      fieldnames=fieldnames,
                                      delimiter=',')

            self.csv.writeheader()

            for groupName in self._traverseGroups():
                for user in self._traverseUsersWithinGroup(groupName):
                    self.csv.writerow(user)

            csvFile.close()

    def storeGroupsToFile(self):
        fieldnames = ['groupName']
        if pargs.usersFile:
            fileName = pargs.groupsFile
        else:
            fileName = pargs.userPoolId + '-groups.csv'

        with open(fileName, 'w', newline='') as csvFile:
            self.csv = csv.DictWriter(csvFile,
                                      fieldnames=fieldnames,
                                      delimiter=',')

            self.csv.writeheader()

            for groupName in self._traverseGroups():
                self.csv.writerow({'groupName': groupName})

            csvFile.close()

    def importUsersToCognito(self):
        with open(pargs.usersFile, newline='') as csvfile:
            users = csv.reader(csvfile, delimiter=',')

            header = []
            run = 0
            for user in users:
                if run == 0:
                    run += 1
                    header = user
                else:
                    userAttrs = dict(zip(header, user))
                    userName = userAttrs.pop('cognito:username')
                    if 'cognito:mfa_enabled' in userAttrs:
                        userAttrs.pop('cognito:mfa_enabled')

                    attr = [{'Name': _, 'Value': userAttrs[_]}
                            for _ in userAttrs]
                    try:
                        self.boto3.admin_create_user(
                            UserPoolId=pargs.userPoolId,
                            Username=userName,
                            UserAttributes=attr)
                    except Exception as e:
                        if 'UsernameExistsException' in str(e):
                            logger.debug(str(e) + ' (userName: ' + userName + ')')
                        elif 'InvalidParameterException' in str(e):
                            logger.debug(str(e) + ' (userName: ' + userName + ')')
                        else:
                            logger.error(str(e) + ' (userName: ' + userName + ')')

    def importGroupsToCognito(self):
        with open(pargs.groupsFile, newline='') as csvfile:
            groups = csv.reader(csvfile, delimiter=',')

            run = 0
            for group in groups:
                params = dict(zip(['GroupName'], group))
                params['UserPoolId'] = pargs.userPoolId
                if run != 0:
                    try:
                        self.boto3.create_group(**params)
                    except Exception as e:
                        if 'GroupExistsException' in str(e):
                            logger.debug(str(e) + ' (groupName: ' + params['GroupName'] + ')')
                        else:
                            logger.error(str(e) + ' (groupName: ' + params['GroupName'] + ')')
                        
                run += 1

    def assignUsersToGroupCognito(self):
        with open(pargs.usersGroupsFile, newline='') as csvfile:
            maps = csv.reader(csvfile, delimiter=',')

            run = 0
            for _map in maps:
                params = dict(zip(['GroupName', 'Username'], _map))
                params['UserPoolId'] = pargs.userPoolId
                if run == 0:
                    run += 1
                else:
                    self.boto3.admin_add_user_to_group(**params)


def main():
    if pargs.exports and pargs.imports:
        logger.error('Only import or export can be done at one time.')
        exit()

    if pargs.exports:
        c = AwsCognito()
        c.storeUsersToFile()
        c.storeGroupsToFile()
        c.storeUsersGroupsMapToFile()

    if pargs.imports:
        users = pargs.usersFile is not None
        groups = pargs.groupsFile is not None
        maps = pargs.usersGroupsFile is not None

        if not users and not groups and not maps:
            logger.error('At least one of -f, -g or -m switch must be givem ' +
                         'if import is considered')
        else:
            c = AwsCognito()

        if users:
            c.importUsersToCognito()
            pass

        if groups:
            c.importGroupsToCognito()
            pass

        if maps:
            c.assignUsersToGroupCognito()
            pass


if __name__ == '__main__':
    main()
