# © 2022 Amazon Web Services, Inc. or its affiliates. All Rights Reserved.
#
# SimpleIOT project.
# Author: Ramin Firoozye (framin@amazon.com)
#
# These are common configuration management routines used by subsystems that need
# access to configuration data.
#
# Note that a super-set of this is in the iotcli folder so it can be run standalone
# with a minimal version of the config.json file for those who don't need to
# download the whole package. Also, note that this used to be a symbolic link to that
# file, but after the refactoring, we copied it from there so there would be no
# external dependencies to the installer, but the two files have diverged since.
# If changes are made to that file, this one has to be manually merged.
#
# The other config file is at: simpleiot/sources/iotapi/util/config.py
#
# The configuration system for SimpleIOT is based on files placed in the following
# hierarchy:
#
# ~ (home directory)
#  .simpleiot (root of SimmpleIOT)
#     {team-name}
#        bootstrap.json - json file generated by the bootstrap process
#        cdkoutput.json - json file output from running the CDK
#        config.json - merged JSON file input by other subsystems
#        projects
#           {project-name}
#           ...
#                models
#                   {model-name}
#                   ...
#                       {devices}
#                          {device-1-serial}
#                          {device-2}
#                              {serial}-cert.pem
#                              {serial}-private-key.pem
#                              {serial}-public-key.pem
#
import re
import shutil
import stat
import os
import json
import tempfile
import boto3
from botocore.exceptions import ClientError
import questionary
import keyring
import webbrowser
import traceback
from questionary import Validator, ValidationError, prompt
import click
import functools


SIMPLEIOT_LOCAL_ROOT = "~/.simpleiot"
DEFAULT_TEAM = "simpleiot"


# This just makes sure the directory structure is set up properly, and has proper access
# privileges. The 'sync' mechanism re-creates all the certs with the devices and models.
#
# NOTE also that the serial numbers MUST be normalized so they can be supported by the
# local filesystem.
#
# NOTE: if the user has specified a different project root, we still look for it in the
# config.json file in ~/.iot.
#

class Config(object):
    def __init__(self, team, use_sso, profile, account, region, identity, userpool,
                 client_id, api_endpoint, config, debug):
        self.team = team
        self.use_sso = use_sso
        self.profile = profile
        self.account = account
        self.region = region
        self.identity = identity
        self.userpool = userpool
        self.client_id = client_id
        self.config = config
        self.debug = debug
        self.api_endpoint = api_endpoint

        # We get any stored API tokens from the keyring
        # If not defined, we'll have to signal that they have to first do an
        # 'iot auth login --usenname .... --password ...' to get a token.
        #
        # If in SSO mode, we won't be having any such token.
        #
        self.token = get_stored_api_token(self)
        # if not self.token:
        #     print(f"WARNING: missing authentication. Make sure you run 'iot auth login ...' before proceeding")

def preload_config(team,
              profile=None,
              debug=False):

    config = None
    try:
        # config_file = None
        # conf_dir = get_conf_dir()
        # if not conf_dir:
        #     if profile == None:
        #         config_path = os.path.join(conf_dir, "config.json")
        #     else:
        #         config_path = os.path.join(conf_dir, f"config-{profile}.json")
        #
        # with open(config_path, "r") as infile:
        #     config = json.load(infile)

        config_data = load_config(team)
        if config_data:
            # NOTE: the camelcase parameters are generated by CDK/CloudFormation.
            # CloudFormation does not allow underlines in output names.
            #
            team = config_data.get("team", team)
            account = config_data.get("account", None)
            region = config_data.get("region", None)
            api_endpoint = config_data.get("apiEndpoint", None)
            identity = config_data.get("cognitoIdentityPoolId", None)
            client_id = config_data.get("cognitoClientId", None)

            use_sso = config_data.get("use_sso", False)

            if use_sso:
                profile = None
                userpool = None
            else:
                profile = config_data.get("aws_profile", None)
                userpool = config_data.get("cognitoUserPoolId", None)

            if debug:
                if team:
                    print(f"Team ID: {team}")
                if profile:
                    print(f"Profile: {profile}")
                if account:
                    print(f"Account: {account}")
                if region:
                    print(f"Region: {region}")
                if identity:
                    print(f"Cognito Identity Pool: {identity}")
                if userpool:
                    print(f"Cognito User Pool: {userpool}")
                if client_id:
                    print(f"Cognito Client ID: {client_id}")
                if api_endpoint:
                    print(f"Endpoint: {api_endpoint}")

            config = Config(team, use_sso, profile, account, region, identity,
                            userpool, client_id, api_endpoint, config, debug)
        return config
    except Exception as e:
        print(f"ERROR: could not open configuration file for profile in ~/.simpleiot. {str(e)}")
        exit(1)



# This is a common wrapper to be used by all commands to add common parameters. Note that
# adding these means the options MUST be on the function definition.

def common_cli_params(func):
    @click.option("--team", "-t", envvar="IOT_TEAM", help="IOT Team ID", required=True)
    @click.option("--profile", "-pr", help="AWS Profile", envvar="AWS_PROFILE", default=None)
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


def _normalize_path_name(src):
    """
    Normalize a path name by removing all characters that aren't in a sanitized set.
    :param src: source name
    :return: cleaned path name
    """
    clean = re.sub('[^a-zA-Z0-9_-]', '', src)
    return clean


def list_all_teams():
    """
    Return a list of all teams. These are identified as directories inside the ~/.simpleiot
    directory that contain a config.json file. We ignore directories that start with an underline.
    These are for internal use.
    """
    result = {}
    root_iot_dir = os.path.expanduser(SIMPLEIOT_LOCAL_ROOT)
    subfolders = [f.name for f in os.scandir(root_iot_dir) if f.is_dir() and f.name[0] != "_"]

    for team_name in subfolders:
        try:
            config = load_config(team_name, exit_on_error=False)
            if config:
                result[team_name] = config
        except Exception as e:
            pass

    return result

def get_iot_team_dir(team=DEFAULT_TEAM, create=True):
    """
    At the root of the IOT directory, we create directories for each
    given team. Inside each one will be the subdirectories for projects,
    models, etc. The team name will be 'normalized' by taking out
    invalid characters.

    Note that this may cause collisions if someone creates team names
    that differ only with invalid characters, for example ABC$ and ABC!
    will both be normalized to ABC.

    :param team: name of team.
    :param create: If set to true, creates the directory path if it doesn't exist
    :return:
    """
    result = None

    try:
        root_iot_dir = os.path.expanduser(SIMPLEIOT_LOCAL_ROOT)
        if not os.path.exists(root_iot_dir):
            if create:
                os.makedirs(root_iot_dir)

        normal_team = _normalize_path_name(team)
        team_dir = os.path.join(root_iot_dir, normal_team)
        if not os.path.exists(team_dir):
            if create:
                os.makedirs(team_dir)

        result = team_dir
    except Exception as e:
        print("ERROR creating IOT root directory: " + str(e))

    return result


def get_iot_project_dir(team, project, create=True):
    """
    This returns the root of a single project directory.

    :param team: Team name
    :param project: Project name
    :param create: If set to true, creates the directory path if it doesn't exist
    :return:
    """
    result = None
    team_root = get_iot_team_dir(team, create)
    if team_root:
        normal_project = _normalize_path_name(project)
        project_dir = os.path.join(team_root, "projects", normal_project)
        if create:
            if not os.path.exists(project_dir):
                os.makedirs(project_dir)

        result = project_dir
    return result


def get_iot_model_dir(team, project, model, create=True):
    """
    This returns the root directory where model information is saved.

    :param team: Team name
    :param project: Project name
    :param model: Model name
    :param create: If set to true, creates the directory path if it doesn't exist
    :return:
    """
    result = None
    project_dir = get_iot_project_dir(team, project, create)
    if project_dir:
        normal_model = _normalize_path_name(model)
        model_dir = os.path.join(project_dir, "models", normal_model)
        if create:
            if not os.path.exists(model_dir):
                os.makedirs(model_dir)

        result = model_dir
    return result


def get_iot_device_dir(team, project, model, device, create=True):
    """
    This returns the root directory where individual device information is saved.

    :param team: Team name
    :param project: Project name
    :param model: Model name
    :param device: Device serial ID
    :param create: If set to true, creates the directory path if it doesn't exist
    :return:
    """
    result = None
    model_dir = get_iot_model_dir(team, project, model, create)
    if model_dir:
        normal_device = _normalize_path_name(device)
        device_dir = os.path.join(model_dir, "devices", normal_device)
        if create:
            if not os.path.exists(device_dir):
                os.makedirs(device_dir)

        result = device_dir
    return result


def delete_even_readonly_files(action, name, exc):
    """
    Utility to force a file to writeable and then remove it.
    This is passed as a shutil.rmtree error handler to clean up a file
    before deletion.

    :param action:
    :param name:
    :param exc:
    :return:
    """
    try:
        os.chmod(name, stat.S_IWRITE)
        os.remove(name)
    except Exception:
        pass


def delete_iot_device_dir(team, project, model, device):
    """
    Delete all the on-disk cache data having to do with a single Device.

    :param team:
    :param project:
    :param model:
    :param device:
    :return:
    """
    device_dir = get_iot_device_dir(team, project, model, device, create=False)
    #print(f"GOT device directory: {device_dir}")
    if device_dir:
        shutil.rmtree(device_dir, onerror=delete_even_readonly_files)


def delete_iot_model_dir(team, project, model):
    """
    Delete all the on-disk cache data having to do with a single Model.

    :param team:
    :param project:
    :param model:
    :return:
    """
    model_dir = get_iot_model_dir(team, project, model, create=False)
    if model_dir:
        shutil.rmtree(model_dir, onerror=delete_even_readonly_files)


def delete_iot_project_dir(team, project):
    """
    Delete all the on-disk cache data having to do with a single Project.

    :param team:
    :param project:
    :return:
    """
    project_dir = get_iot_project_dir(team, project, create=False)
    if project_dir:
        shutil.rmtree(project_dir, onerror=delete_even_readonly_files)


def delete_iot_team_dir(team):
    """
    Delete all the on-disk cache data having to do with a single Team.

    :param team:
    :return:
    """
    team_dir = get_iot_team_dir(team, create=False)
    if team_dir:
        shutil.rmtree(team_dir, onerror=delete_even_readonly_files)

#################################
# Configuration loading routines
#
# These are used to centrally load config files for each Team
#
def path_for_bootstrap_file(team=DEFAULT_TEAM):
    team_path = get_iot_team_dir(team)
    bootstrap_path = os.path.join(team_path, "bootstrap.json")
    return bootstrap_path


def path_for_cdkoutput_file(team=DEFAULT_TEAM):
    team_path = get_iot_team_dir(team)
    cdkoutput_path = os.path.join(team_path, "cdkoutput.json")
    return cdkoutput_path

# A config file is a merged JSON file that combines boothstrap and cdkoutput
# data.

def path_for_config_file(team=DEFAULT_TEAM):
    team_path = get_iot_team_dir(team)
    config_path = os.path.join(team_path, "config.json")
    return config_path


def path_for_certs(team=DEFAULT_TEAM):
    team_path = get_iot_team_dir(team)
    certs_path = os.path.join(team_path, "certs")
    if not os.path.exists(certs_path):
        os.makedirs(certs_path)
    result = certs_path
    return result


def load_bootstrap_config(team=DEFAULT_TEAM):
    config_data = None
    bootstrap_path = None
    try:
        bootstrap_path = path_for_bootstrap_file(team)

        with open(bootstrap_path, "r") as infile:
            config_data = json.load(infile)
    except Exception as e:
        print(f"ERROR loading bootstrap config file: [{bootstrap_path}]: {str(e)}")
        exit(1)

    return config_data


def load_defaults_config(defaults_path):
    config_data = None
    try:
        with open(defaults_path, "r") as infile:
            config_data = json.load(infile)
    except Exception as e:
        print(f"ERROR loading defaults config file: [{defaults_path}]: {str(e)}")
        exit(1)

    return config_data


def load_cdkoutput_config(team=DEFAULT_TEAM):
    config_data = None
    cdkoutput_path = None
    try:
        cdkoutput_path = path_for_cdkoutput_file(team)

        with open(cdkoutput_path, "r") as infile:
            config_root = json.load(infile)
            config_data = config_root.get("Iotcdk", None)
    except Exception as e:
        print(f"ERROR loading CDK output config file: [{cdkoutput_path}]: {str(e)}")
        exit(1)

    return config_data


def does_team_exist(team=DEFAULT_TEAM):
    result = False
    config_data = None
    config_path = None
    try:
        config_path = path_for_config_file(team)

        with open(config_path, "r") as infile:
            config_data = json.load(infile)

        account = config_data.get("account")
        result = True

    except Exception as e:
        pass

    return result


def load_config(team=DEFAULT_TEAM, exit_on_error=True):
    config_data = None
    config_path = None
    try:
        config_path = path_for_config_file(team)
        if not os.path.exists(config_path):
            if exit_on_error:
                print(f"ERROR: team '{team}' not recognized.")
                exit(1)
            else:
                return config_data # None

        with open(config_path, "r") as infile:
            config_data = json.load(infile)
    except Exception as e:
        print(f"ERROR: could not locate configuration data for project [{team}].")

    return config_data


####################################################################################
# Temp files - these are used to pass data between various phases of installation.
#
def save_to_tempfile(filename, data):
    try:
        tempdir = tempfile.gettempdir()
        outfile = os.path.join(tempdir, filename)
        with open(outfile, 'w') as out:
            out.write(data + '\n')
    except Exception as e:
        print(f"ERROR: could not save data to tempfile: {str(e)}")
        raise e


def load_from_tempfile(filename):
    try:
        tempdir = tempfile.gettempdir()
        infile = os.path.join(tempdir, filename)
        with open(infile, 'r') as out:
            result = out.read().replace('\n','')
            return result
    except Exception as e:
        print(f"ERROR: could not read data from tempfile {filename}")
        raise e


def load_from_tempfile_and_delete(filename):
    try:
        tempdir = tempfile.gettempdir()
        infile = os.path.join(tempdir, filename)
        with open(infile, 'r') as out:
            result = out.read().replace('\n','')
        os.remove(infile)
        return result
    except Exception as e:
        print(f"ERROR: could not read data from tempfile {filename}")
        raise e


#######

def ask_to_confirm_delete():
    is_delete = questionary.text("\nAre you sure you want to do this. Enter 'DELETE' to confirm:").ask()
    if is_delete == 'DELETE':
        return True
    else:
        return False


def ask_to_confirm_yesno(prompt):
    response = False
    confirm = questionary.text(f"{prompt} [N/y]:").ask()
    if confirm:
        confirm_l = confirm.lower()
        if confirm_l == 'y' or confirm_l == "ye" or confirm_l == "yes":
           response = True

    return response

#########

#
# Username/password/API token/SSO session management
#
# NOTE: a service name is synthesized with the team name. If not specified
# in the config file, we use default. This way we can store tokens and passwords
# for multiple teams.
#
def _get_service(config):
    team = config.team
    if not team:
        team = DEFAULT_TEAM

    service = f"simpleiot_{team}"
    return service

def get_stored_key(config, key):
    service = _get_service(config)
    value = keyring.get_password(service, key)
    return value

def set_stored_key(config, key, value):
    service = _get_service(config)
    keyring.set_password(service, key, value)

def get_stored_username(config):
    return get_stored_key(config, "username")

def get_stored_password(config):
    return get_stored_key(config, "password")

def get_stored_api_token(config):
    return get_stored_key(config, "api_token")

def get_stored_access_key(config):
    return get_stored_key(config, "access_key")

def get_stored_access_secret(config):
    return get_stored_key(config, "access_secret")

def get_stored_session_token(config):
    return get_stored_key(config, "session_token")

def store_username(config, username):
    set_stored_key(config, "username", username)

def store_password(config, password):
    set_stored_key(config, "password", password)

def store_api_token(config, token):
    set_stored_key(config, "api_token", token)

def store_access_key(config, key):
    set_stored_key(config, "access_key", key)

def store_access_secret(config, secret):
    set_stored_key(config, "access_secret", secret)

def store_session_token(config, token):
    set_stored_key(config, "session_token", token)

def clear_api_token(config):
    try:
        service = _get_service(config)
        keyring.delete_password(service, "api_token")
    except Exception as e:
        pass

def clear_sso_tokens(config):
    try:
        service = _get_service(config)
        keyring.delete_password(service, "access_key")
        keyring.delete_password(service, "access_secret")
        keyring.delete_password(service, "session_token")
    except Exception as e:
        pass

def clear_all_auth(config):
    try:
        service = _get_service(config)
        keyring.delete_password(service, "username")
        keyring.delete_password(service, "password")
        keyring.delete_password(service, "api_token")
        keyring.delete_password(service, "access_key")
        keyring.delete_password(service, "access_secret")
        keyring.delete_password(service, "session_token")
    except Exception as e:
        pass

################
# SSO Login

class IsEmptyValidator(Validator):
    def validate(self, document):
        if len(document.text) == 0:
            raise ValidationError(
                message="Please enter a value",
                cursor_position=len(document.text),
            )


#
# This function tries to login using SSO credentials. If SSO credentials were saved from the last time
# they're used here. If not, we ask for them.
#
def login_with_sso(config):
    try:
        GRANT_TYPE = 'urn:ietf:params:oauth:grant-type:device_code'
        sso_url_stored = config.get("sso_url", "")
        url_ok = False
        import validators

        while not url_ok:
            sso_url = questionary.text("AWS SSO url?", validate=IsEmptyValidator, default=sso_url_stored).ask()
            url_ok = validators.url(sso_url)
            if not url_ok:
                print("ERROR: Invalid URL")

        sso_region_stored = config.get("region", None)
        session = boto3.session.Session()
        regions = session.get_available_regions('sso-oidc')
        if len(regions) == 0:
            print("ERROR: SSO not supported in these regions")
            exit(1)

        if not sso_region_stored:
            sso_region_stored = regions[0]

        sso_region = questionary.select("AWS SSO region: ", choices=regions, default=sso_region_stored).ask()

        sso_oidc = boto3.client('sso-oidc', region_name=sso_region)
        try:
            resp = sso_oidc.register_client(clientName="simpleiot", clientType="public")
            client_id = resp['clientId']
            client_secret = resp['clientSecret']

            resp = sso_oidc.start_device_authorization(
                clientId=client_id, clientSecret=client_secret, startUrl=sso_url)
            device_code = resp['deviceCode']
            signin_url = resp['verificationUriComplete']
            webbrowser.open(signin_url)

            access_token = False

            while not access_token:
                try:
                    resp = sso_oidc.create_token(clientId=client_id, clientSecret=client_secret,
                                                 grantType=GRANT_TYPE, deviceCode=device_code)
                    access_token = resp['accessToken']
                except sso_oidc.exceptions.AuthorizationPendingException:
                    continue
                except Exception as e:
                    print(f"ERROR: CREATE SSO TOKEN EXCEPTION: {str(e)}")
                    exit(1)

            sso = boto3.client('sso', region_name=sso_region)

            resp = sso.list_accounts(maxResults=1000, accessToken=access_token)
            account_array = resp['accountList']
            account_id = None
            account_name= None
            account_role = None

            if len(account_array) == 0:
                print("ERROR: No valid SSO accounts found")
                exit(1)

            # If we only have one account, we don't bother asking
            #
            if len(account_array) == 1:
                account = account_array[0]
                account_id = account['accountId']
                account_name = account['accountName']
            else:
                account_list = {}
                account_names = {}
                for account in account_array:
                    _name = account['accountName']
                    _email = account['emailAddress']
                    _id = account['accountId']
                    account_list[f"{_name} ({_email})"] = _id
                    account_names[_id] = _name

                account_names = list(account_list.keys())
                account_choice = questionary.select("Select SSO Account: ", choices=account_names).ask()
                account_id = account_list[account_choice]
                account_name = account_names[account_id]

            resp = sso.list_account_roles(maxResults=1000,
                                          accessToken=access_token,
                                          accountId=account_id)
            role_array = resp['roleList']
            if len(role_array) == 0:
                print("ERROR: no roles found for account. Please correct and re-run")
                exit(1)

            if len(role_array) == 1:
                account_role = role_array[0]['roleName']
            else:
                role_list = []
                for role in role_array:
                    role_list.append(role['roleName'])

                account_role = questionary.select("Select SSO Account Role: ", choices=role_list).ask()

            resp = sso.get_role_credentials(roleName=account_role, accountId=account_id,
                                            accessToken=access_token)
            access_key = resp['roleCredentials']['accessKeyId']
            access_secret = resp['roleCredentials']['secretAccessKey']
            session_token = resp['roleCredentials']['sessionToken']

            return sso_url, \
                   sso_region, \
                   account_id, \
                   account_role, \
                   account_name, \
                   access_key, \
                   access_secret, \
                   session_token

        except sso_oidc.exceptions.InvalidRequestException:
            print("ERROR: Invalid Request to SSO")
            exit(1)
        except sso_oidc.exceptions.TooManyRequestsException:
            print("ERROR: Too many requests Exception")
            exit(1)
        except sso_oidc.exceptions.UnauthorizedException:
            print("ERROR: Unauthorized Exception")
            exit(1)
        except sso_oidc.exceptions.ResourceNotFoundException:
            print("ERROR: Resource Not Found Exception")
            exit(1)
        except Exception as e:
            print(f"ERROR: {str(e)}")
            traceback.print_exc()
            return None

    except Exception as e:
        print(f"ERROR: {str(e)}")
        traceback.print_exc()
        return None

##############################################################################
# This allows looking up a string from elements of an Enum.
# Used to lookup device model fields. The lookup is case-insensitive.
#
def enum_from_str(en, nm):
    try:
        for item in en:
            if nm.lower() == item.name.lower():
                return item.value

        raise ValueError(f"Invalid value specified for enum {str(en)}")
    except Exception as e:
        print(f"Error: {e}")
        return None

#
# This reverses the above and returns the string value of an enum.
# Used to display or export an enum into a config file.
#
def enum_to_str(en, item):
    result = None
    try:
        result = en(item).name.lower()
    except Exception as e:
        pass

    return result

# This does the above, except it looks for a list of string values,
# which it then turns into an ORed bitflag.
#
# The values are assumed to be comma-separated and each a valid item in the
# enum. We don't check for duplicates.
#
def enum_from_str_list(en, nml):
    flag = 0
    slist = nml.split(",")
    try:
        for st in slist:
            for item in en:
                if st.strip().lower() == item.name.lower():
                    flag |= int(item.value)
                    break
        return flag
    except Exception as e:
        print(f"Error: {e}")
        return None

#
# This does the reverse, except it returns a list of enums.
#
def enum_to_str_list(en, value):
    result = []
    try:
        for b in range(1, len(en) + 1):
            cb = b - 1
            item = en(1 << cb)
            if (value & (1 << (b - 1))):
                result.append(item.name)
        return ", ".join(result)
    except Exception as e:
        print(f"ERROR: {str(e)}")
        pass
    return result

#
# This is used to get a secret out of secretsmanager.
#
def get_secret(config, name):
    profile = config.get("aws_profile", "default")
    os.environ['AWS_PROFILE'] = profile
    session = boto3.session.Session()
    client = session.client(
        service_name='secretsmanager',
        region_name=config.get("region", None)
    )

    # In this sample we only handle the specific exceptions for the 'GetSecretValue' API.
    # See https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
    # We rethrow the exception by default.

    try:
        secret_value = client.get_secret_value(
            SecretId=name
        )
    except ClientError as e:
        print(f"Got exception: {str(e)}")

        if e.response['Error']['Code'] == 'DecryptionFailureException':
            # Secrets Manager can't decrypt the protected secret text using the provided KMS key.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'InternalServiceErrorException':
            # An error occurred on the server side.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'InvalidParameterException':
            # You provided an invalid value for a parameter.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'InvalidRequestException':
            # You provided a parameter value that is not valid for the current state of the resource.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'ResourceNotFoundException':
            # We can't find the resource that you asked for.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
    else:
        # Decrypts secret using the associated KMS CMK.
        # Depending on whether the secret is a string or binary, one of these fields will be populated.
        if 'SecretString' in secret_value:
            secret = secret_value["SecretString"]
            return json.loads(secret)
        else:
            return None
