from __future__ import unicode_literals, print_function
from builtins import chr
import base64
import json
import os
import pickle

from Crypto import Random
from Crypto.Cipher import AES

from cumulusci.core.config import BaseConfig
from cumulusci.core.config import ConnectedAppOAuthConfig
from cumulusci.core.config import OrgConfig
from cumulusci.core.config import ScratchOrgConfig
from cumulusci.core.config import ServiceConfig
from cumulusci.core.exceptions import OrgNotFound
from cumulusci.core.exceptions import ServiceNotConfigured
from cumulusci.core.exceptions import ServiceNotValid
from cumulusci.core.exceptions import KeychainConnectedAppNotFound


class BaseProjectKeychain(BaseConfig):
    encrypted = False

    def __init__(self, project_config, key):
        super(BaseProjectKeychain, self).__init__()
        self.config = {
            'orgs': {},
            'app': None,
            'services': {},
        }
        self.project_config = project_config
        self.key = key
        self._load_keychain()

    def _load_keychain(self):
        self._load_app()
        self._load_orgs()
        self._load_scratch_orgs()
        self._load_services()

    def _load_app(self):
        pass

    def _load_orgs(self):
        pass

    def _load_scratch_orgs(self):
        """ Creates all scratch org configs for the project in the keychain if
            a keychain org doesn't already exist """
        current_orgs = self.list_orgs()
        if not self.project_config.orgs__scratch:
            return
        for config_name in self.project_config.orgs__scratch.keys():
            if config_name in current_orgs:
                # Don't overwrite an existing keychain org
                continue
            self.create_scratch_org(config_name, config_name)

    def _load_services(self):
        pass
            
    def create_scratch_org(self, org_name, config_name, days=None):
        """ Adds/Updates a scratch org config to the keychain from a named config """
        scratch_config = getattr(self.project_config, 'orgs__scratch__{}'.format(config_name))
        if days is not None:
            # Allow override of scratch config's default days
            scratch_config['days'] = days
        else:
            # Use scratch config days or default of 1 day
            scratch_config.setdefault('days', 1)
        scratch_config['scratch'] = True
        scratch_config.setdefault('namespaced', False)
        scratch_config['config_name'] = config_name
        scratch_config['sfdx_alias'] = '{}__{}'.format(
            self.project_config.project__name,
            org_name,
        )
        org_config = ScratchOrgConfig(scratch_config, org_name)
        self.set_org(org_config)

    def change_key(self, key):
        """ re-encrypt stored services, orgs, and the connected_app
        with the new key """

        connected_app = self.get_connected_app()

        services = {}
        for service_name in self.list_services():
            services[service_name] = self.get_service(service_name)

        orgs = {}
        for org_name in self.list_orgs():
            orgs[org_name] = self.get_org(org_name)

        self.key = key
        if connected_app:
            self.set_connected_app(connected_app)

        if orgs:
            for org_name, org_config in list(orgs.items()):
                self.set_org(org_config)

        if services:
            for service_name, service_config in list(services.items()):
                self.set_service(service_name, service_config)

    def set_connected_app(self, app_config, project=False):
        """ store a connected_app configuration """

        self._set_connected_app(app_config, project)
        self._load_app()

    def _set_connected_app(self, app_config, project):
        self.app = app_config

    def get_connected_app(self):
        """ retrieve the connected app configuration """

        return self._get_connected_app()

    def _get_connected_app(self):
        return self.app

    def remove_org(self, name, global_org=None):
        if name in self.orgs.keys():
            self._remove_org(name, global_org)

    def _remove_org(self, name, global_org):
        del self.orgs[name]
        self._load_orgs()

    def set_org(self, org_config, global_org=False):
        if isinstance(org_config, ScratchOrgConfig):
            org_config.config['scratch'] = True
        self._set_org(org_config, global_org)
        self._load_orgs()

    def _set_org(self, org_config, global_org):
        self.orgs[org_config.name] = org_config

    def get_default_org(self):
        """ retrieve the name and configuration of the default org """
        for org in self.list_orgs():
            org_config = self.get_org(org)
            if org_config.default:
                return org, org_config
        return None, None

    def set_default_org(self, name):
        """ set the default org for tasks by name key """
        org = self.get_org(name)
        self.unset_default_org()
        org.config['default'] = True
        self.set_org(org)

    def unset_default_org(self):
        """ unset the default orgs for tasks """
        for org in self.list_orgs():
            org_config = self.get_org(org)
            if org_config.default:
                del org_config.config['default']
                self.set_org(org_config)

    def get_org(self, name):
        """ retrieve an org configuration by name key """
        if name not in self.orgs:
            self._raise_org_not_found(name)
        return self._get_org(name)

    def _get_org(self, name):
        return self.orgs.get(name)

    def _raise_org_not_found(self, name):
        raise OrgNotFound(
            'Org named {} was not found in keychain'.format(name))

    def list_orgs(self):
        """ list the orgs configured in the keychain """
        orgs = list(self.orgs.keys())
        orgs.sort()
        return orgs

    def set_service(self, name, service_config, project=False):
        """ Store a ServiceConfig in the keychain """
        if not self.project_config.services or name not in self.project_config.services:
            self._raise_service_not_valid(name)
        self._validate_service(name, service_config)
        self._set_service(name, service_config, project)
        self._load_services()

    def _set_service(self, name, service_config, project=False):
        self.services[name] = service_config

    def get_service(self, name):
        """ Retrieve a stored ServiceConfig from the keychain or exception

        :param name: the service name to retrieve
        :type name: str

        :rtype ServiceConfig
        :return the configured Service
        """
        if not self.project_config.services or name not in self.project_config.services:
            self._raise_service_not_valid(name)
        if name not in self.services:
            self._raise_service_not_configured(name)

        return self._get_service(name)

    def _get_service(self, name):
        return self.services.get(name)

    def _validate_service(self, name, service_config):
        missing_required = []
        attr_key = 'services__{0}__attributes'.format(name)
        for atr, config in list(getattr(self.project_config, attr_key).items()):
            if config.get('required') is True and not getattr(service_config, atr):
                missing_required.append(atr)

        if missing_required:
            self._raise_service_not_valid(name)

    def _raise_service_not_configured(self, name):
        raise ServiceNotConfigured(
            'Service named {} is not configured for this project. Configured services are: {}'.format(name, ', '.join(list(self.services)))
        )

    def _raise_service_not_valid(self, name):
        raise ServiceNotValid(
            'Service named {} is not valid for this project'.format(name))

    def list_services(self):
        """ list the services configured in the keychain """
        services = list(self.services.keys())
        services.sort()
        return services


class EnvironmentProjectKeychain(BaseProjectKeychain):
    """ A project keychain that stores org credentials in environment variables """
    encrypted = False
    org_var_prefix = 'CUMULUSCI_ORG_'
    app_var = 'CUMULUSCI_CONNECTED_APP'
    service_var_prefix = 'CUMULUSCI_SERVICE_'

    def _get_env(self):
        """ loads the environment variables as unicode if ascii """
        try:
            return [(k.decode(), v.decode()) for k,v in list(os.environ.items())]
        except AttributeError:
            return list(os.environ.items())

    def _load_app(self):
        try:
            app = os.environ.get(self.app_var.encode('ascii'))
        except TypeError:
            app = os.environ.get(self.app_var)
        if app:
            self.app = ConnectedAppOAuthConfig(json.loads(app))

    def _load_orgs(self):
        for key, value in self._get_env():
            if key.startswith(self.org_var_prefix):
                org_config = json.loads(value)
                org_name = key[len(self.org_var_prefix):]
                if org_config.get('scratch'):
                    self.orgs[org_name] = ScratchOrgConfig(json.loads(value), org_name)
                else:
                    self.orgs[org_name] = OrgConfig(json.loads(value), org_name)

    def _load_services(self):
        for key, value in self._get_env():
            if key.startswith(self.service_var_prefix):
                service_config = json.loads(value)
                service_name = key[len(self.service_var_prefix):]
                self._set_service(service_name, ServiceConfig(service_config))



BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS).encode('ascii')
unpad = lambda s: s[0:-ord(s[-1])]


class BaseEncryptedProjectKeychain(BaseProjectKeychain):
    """ Base class for building project keychains that use AES encryption for securing stored org credentials """
    encrypted = True

    def _set_connected_app(self, app_config, project):
        encrypted = self._encrypt_config(app_config)
        self._set_encrypted_connected_app(encrypted, project)

    def _set_encrypted_connected_app(self, encrypted, project):
        self.app = encrypted

    def _get_connected_app(self):
        if self.app:
            return self._decrypt_config(ConnectedAppOAuthConfig, self.app)

    def _get_service(self, name):
        return self._decrypt_config(ServiceConfig, self.services[name])

    def _set_service(self, service, service_config, project):
        encrypted = self._encrypt_config(service_config)
        self._set_encrypted_service(service, encrypted, project)

    def _set_encrypted_service(self, service, encrypted, project):
        self.services[service] = encrypted

    def _set_org(self, org_config, global_org):
        encrypted = self._encrypt_config(org_config)
        self._set_encrypted_org(org_config.name, encrypted, global_org)

    def _set_encrypted_org(self, name, encrypted, global_org):
        self.orgs[name] = encrypted

    def _get_org(self, name):
        return self._decrypt_config(OrgConfig, self.orgs[name], extra=[name])

    def _get_cipher(self, iv=None):
        if iv is None:
            iv = Random.new().read(AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return cipher, iv

    def _encrypt_config(self, config):
        pickled = pickle.dumps(config.config)
        pickled = pad(pickled)
        #pickled = base64.b64encode(pickled)
        cipher, iv = self._get_cipher()
        encrypted = base64.b64encode(iv + cipher.encrypt(pickled))
        return encrypted

    def _decrypt_config(self, config_class, encrypted_config, extra=None):
        if not encrypted_config:
            if extra:
                return config_class(None, *extra)
            else:
                return config_class()
        encrypted_config = base64.b64decode(encrypted_config)
        iv = encrypted_config[:16]
        cipher, iv = self._get_cipher(iv)
        pickled = cipher.decrypt(encrypted_config[16:])
        config_dict = pickle.loads(pickled)
        args = [config_dict]
        if extra:
            args += extra
        return self._construct_config(config_class, args)

    def _construct_config(self, config_class, args):
        if args[0].get('scratch'):
            config_class = ScratchOrgConfig
        
        return config_class(*args)


class EncryptedFileProjectKeychain(BaseEncryptedProjectKeychain):
    """ An encrypted project keychain that stores in the project's local directory """

    @property
    def config_local_dir(self):
        return os.path.join(
            os.path.expanduser('~'),
            self.project_config.global_config_obj.config_local_dir,
        )

    @property
    def project_local_dir(self):
        return self.project_config.project_local_dir

    def _load_files(self, dirname, extension, key):
        for item in os.listdir(dirname):
            if item.endswith(extension):
                with open(os.path.join(dirname, item), 'r') as f_item:
                    config = f_item.read()
                name = item.replace(extension, '')
                if not key in self.config:
                    self.config[key] = []
                self.config[key][name] = config

    def _load_file(self, dirname, filename, key):
        full_path = os.path.join(dirname, filename)
        if not os.path.exists(full_path):
            return
        with open(os.path.join(dirname, filename), 'r') as f_item:
            config = f_item.read()
        self.config[key] = config

    def _load_app(self):
        self._load_file(self.config_local_dir, 'connected.app', 'app')
        self._load_file(self.project_local_dir, 'connected.app', 'app')

    def _load_orgs(self):
        self._load_files(self.config_local_dir, '.org', 'orgs')
        self._load_files(self.project_local_dir, '.org', 'orgs')

    def _load_services(self):
        self._load_files(self.config_local_dir, '.service', 'services')
        self._load_files(self.project_local_dir, '.service', 'services')

    def _remove_org(self, name, global_org):
        if global_org:
            full_path = os.path.join(self.config_local_dir, '{}.org'.format(name))
        else:
            full_path = os.path.join(self.project_local_dir, '{}.org'.format(name))
        if not os.path.exists(full_path):
            kwargs = {'name': name}
            if not global_org:
                raise OrgNotFound('Could not find org named {name} to delete.  Deleting in project org mode.  Is {name} a global org?'.format(**kwargs))
            raise OrgNotFound('Could not find org named {name} to delete.  Deleting in global org mode.  Is {name} a project org instead of a global org?'.format(**kwargs))
           
        os.remove(full_path) 
        self._load_orgs()

    def _set_encrypted_connected_app(self, encrypted, project):
        if project:
            filename = os.path.join(self.project_local_dir, 'connected.app')
        else:
            filename = os.path.join(self.config_local_dir, 'connected.app')
        with open(filename, 'wb') as f_org:
            f_org.write(encrypted)
        self.app = encrypted

    def _set_encrypted_org(self, name, encrypted, global_org):
        if global_org:
            filename = os.path.join(
                self.config_local_dir, '{}.org'.format(name))
        else:
            filename = os.path.join(
                self.project_local_dir, '{}.org'.format(name))
        with open(filename, 'wb') as f_org:
            f_org.write(encrypted)

    def _set_encrypted_service(self, name, encrypted, project):
        if project:
            filename = os.path.join(
                self.project_local_dir, '{}.service'.format(name))
        else:
            filename = os.path.join(
                self.config_local_dir, '{}.service'.format(name))
        with open(filename, 'wb') as f_service:
            f_service.write(encrypted)

    def _raise_org_not_found(self, name):
        raise OrgNotFound(
            'Org information could not be found.  Expected to find encrypted file at {}/{}.org'.format(
                self.project_local_dir,
                name
            )
        )

    def _raise_service_not_configured(self, name):
        raise ServiceNotConfigured(
            'Service configuration could not be found.  Expected to find encrypted file at {}/{}.org'.format(
                self.project_local_dir,
                name
            )
        )
