from __future__ import print_function, division, absolute_import

try:
    import ssl
except ImportError:
    ssl = None

import dask


_roles = ['client', 'scheduler', 'worker']

_tls_per_role_fields = ['key', 'cert']

_tls_fields = ['ca_file', 'ciphers']

_misc_fields = ['require_encryption']

_fields = set(_misc_fields +
              ['tls_%s' % field for field in _tls_fields] +
              ['tls_%s_%s' % (role, field)
               for role in _roles
               for field in _tls_per_role_fields]
              )


def _field_to_config_key(field):
    return field.replace('_', '-')


class Security(object):
    """
    An object to gather and pass around security configuration.
    Default values are gathered from the global ``config`` object and
    can be overriden by constructor args.

    Supported fields:
        - require_encryption
        - tls_ca_file
        - tls_ciphers
        - tls_client_key
        - tls_client_cert
        - tls_scheduler_key
        - tls_scheduler_cert
        - tls_worker_key
        - tls_worker_cert
    """

    __slots__ = tuple(_fields)

    def __init__(self, **kwargs):
        self._init_from_dict(dask.config.config)
        for k, v in kwargs.items():
            if v is not None:
                setattr(self, k, v)
        for k in _fields:
            if not hasattr(self, k):
                setattr(self, k, None)

    def _init_from_dict(self, d):
        """
        Initialize Security from nested dict.
        """
        self._init_fields_from_dict(d, '', _misc_fields, {})
        self._init_fields_from_dict(d, 'tls', _tls_fields, _tls_per_role_fields)

    def _init_fields_from_dict(self, d, category,
                               fields, per_role_fields):
        if category:
            d = d.get(category, {})
            category_prefix = category + '_'
        else:
            category_prefix = ''
        for field in fields:
            k = _field_to_config_key(field)
            if k in d:
                setattr(self, '%s%s' % (category_prefix, field), d[k])
        for role in _roles:
            dd = d.get(role, {})
            for field in per_role_fields:
                k = _field_to_config_key(field)
                if k in dd:
                    setattr(self, '%s%s_%s' % (category_prefix, role, field), dd[k])

    def __repr__(self):
        items = sorted((k, getattr(self, k)) for k in _fields)
        return ("Security(" +
                ", ".join("%s=%r" % (k, v) for k, v in items if v is not None) +
                ")")

    def get_tls_config_for_role(self, role):
        """
        Return the TLS configuration for the given role, as a flat dict.
        """
        return self._get_config_for_role('tls', role, _tls_fields, _tls_per_role_fields)

    def _get_config_for_role(self, category, role, fields, per_role_fields):
        if role not in _roles:
            raise ValueError("unknown role %r" % (role,))
        d = {}
        for field in fields:
            k = '%s_%s' % (category, field)
            d[field] = getattr(self, k)
        for field in per_role_fields:
            k = '%s_%s_%s' % (category, role, field)
            d[field] = getattr(self, k)
        return d

    def _get_tls_context(self, tls, purpose):
        if tls.get('ca_file') and tls.get('cert'):
            try:
                ctx = ssl.create_default_context(purpose=purpose,
                                                 cafile=tls['ca_file'])
            except AttributeError:
                raise RuntimeError("TLS functionality requires Python 2.7.9+")
            ctx.verify_mode = ssl.CERT_REQUIRED
            # We expect a dedicated CA for the cluster and people using
            # IP addresses rather than hostnames
            ctx.check_hostname = False
            ctx.load_cert_chain(tls['cert'], tls.get('key'))
            if tls.get('ciphers'):
                ctx.set_ciphers(tls.get('ciphers'))
            return ctx

    def get_connection_args(self, role):
        """
        Get the *connection_args* argument for a connect() call with
        the given *role*.
        """
        d = {}
        tls = self.get_tls_config_for_role(role)
        # Ensure backwards compatibility (ssl.Purpose is Python 2.7.9+ only)
        purpose = ssl.Purpose.SERVER_AUTH if hasattr(ssl, "Purpose") else None
        d['ssl_context'] = self._get_tls_context(tls, purpose)
        d['require_encryption'] = self.require_encryption
        return d

    def get_listen_args(self, role):
        """
        Get the *connection_args* argument for a listen() call with
        the given *role*.
        """
        d = {}
        tls = self.get_tls_config_for_role(role)
        # Ensure backwards compatibility (ssl.Purpose is Python 2.7.9+ only)
        purpose = ssl.Purpose.CLIENT_AUTH if hasattr(ssl, "Purpose") else None
        d['ssl_context'] = self._get_tls_context(tls, purpose)
        d['require_encryption'] = self.require_encryption
        return d