django-mellon/mellon/adapters.py

import logging
import uuid
from xml.etree import ElementTree as ET

import lasso
import requests
import requests.exceptions

from django.core.exceptions import PermissionDenied
from django.contrib import auth
from django.contrib.auth.models import Group
from django.utils.text import slugify

from . import utils, app_settings, models
from mellon.federation_utils import idp_metadata_store, url2filename, \
        idp_metadata_extract_entity_id, idp_metadata_is_cached, \
        idp_metadata_load, idp_settings_store, idp_settings_load


class UserCreationError(Exception):
    pass


class DefaultAdapter(object):
    def __init__(self, *args, **kwargs):
        self.logger = logging.getLogger(__name__)

    def get_idp(self, entity_id):
        '''Find the first IdP definition matching entity_id'''
        idp = None
        if idp_metadata_is_cached(entity_id):
            metadata_content = idp_metadata_load(entity_id)
            entity_id = idp_metadata_extract_entity_id(metadata_content)
            idp = {'METADATA': metadata_content,
                   'ENTITY_ID': entity_id}
        else:
            for extra_idp in self.get_identity_providers_setting():
                if extra_idp.get('ENTITY_ID') == entity_id or \
                        idp_metadata_extract_entity_id(extra_idp) == entity_id:
                    idp = extra_idp.copy()

        extra_idp_settings = idp_settings_load(entity_id)
        if extra_idp_settings and idp:
            idp.update(extra_idp_settings)
        return idp

    def get_identity_providers_setting(self):
        for federation_data in self.get_federations():
            if not isinstance(federation_data, dict) or \
                    'FEDERATION' not in federation_data:
                continue
            fed_extra_attrs = federation_data.copy()
            fed_content = fed_extra_attrs.pop('FEDERATION')
            fed_filepath, _ = utils.get_federation_metadata(fed_content)

            try:
                tree = ET.parse(fed_filepath)
                root = tree.getroot()
                for child in root:
                    provider = {}
                    entity_id = idp_metadata_extract_entity_id(ET.tostring(child))
                    if not entity_id:
                        continue
                    provider['METADATA'] = idp_metadata_store(ET.tostring(child))
                    provider.update({'ENTITY_ID': entity_id})
                    provider.update(fed_extra_attrs)
                    idp_settings_store(provider)
                    yield provider
            except:
                self.logger.error('Couldn\'t load federation metadata file %r',
                                  fed_filepath)
                continue

        for extra_provider in app_settings.IDENTITY_PROVIDERS:
            yield extra_provider

    def get_federations(self):
        for federation in getattr(app_settings, 'FEDERATIONS', []):
            yield federation

    def get_idps(self):
        for i, idp in enumerate(self.get_identity_providers_setting()):
            entity_id = idp.get('ENTITY_ID')
            if 'METADATA_URL' in idp and 'METADATA' not in idp:
                verify_ssl_certificate = utils.get_setting(
                    idp, 'VERIFY_SSL_CERTIFICATE')
                try:
                    response = requests.get(idp['METADATA_URL'], verify=verify_ssl_certificate)
                    response.raise_for_status()
                except requests.exceptions.RequestException as e:
                    self.logger.error(
                        u'retrieval of metadata URL %r failed with error %s for %d-th idp',
                        idp['METADATA_URL'], e, i)
                    continue
                md_content = response.content
                if not entity_id:
                    entity_id = idp_metadata_extract_entity_id(md_content)
                idp['METADATA'] = idp_metadata_store(md_content)
            elif not idp.get('METADATA'):
                self.logger.error(u'missing METADATA or METADATA_URL in %d-th idp', i)
                continue
            # load federation-specific configuration
            extra_idp_settings = idp_settings_load(entity_id)
            if extra_idp_settings:
                idp.update(idp_settings_load(entity_id))
            yield idp

    def authorize(self, idp, saml_attributes):
        if not idp:
            return False
        required_classref = utils.get_setting(idp, 'AUTHN_CLASSREF')
        if required_classref:
            given_classref = saml_attributes['authn_context_class_ref']
            if given_classref is None or \
                    given_classref not in required_classref:
                raise PermissionDenied
        return True

    def format_username(self, idp, saml_attributes):
        realm = utils.get_setting(idp, 'REALM')
        username_template = utils.get_setting(idp, 'USERNAME_TEMPLATE')
        try:
            username = unicode(username_template).format(
                realm=realm, attributes=saml_attributes, idp=idp)[:30]
        except ValueError:
            self.logger.error(u'invalid username template %r', username_template)
        except (AttributeError, KeyError, IndexError) as e:
            self.logger.error(
                u'invalid reference in username template %r: %s', username_template, e)
        except Exception as e:
            self.logger.exception(u'unknown error when formatting username')
        else:
            return username

    def create_user(self, user_class):
        return user_class.objects.create(username=uuid.uuid4().hex[:30])

    def finish_create_user(self, idp, saml_attributes, user):
        username = self.format_username(idp, saml_attributes)
        if not username:
            self.logger.warning('could not build a username, login refused')
            raise UserCreationError
        user.username = username
        user.save()

    def lookup_user(self, idp, saml_attributes):
        User = auth.get_user_model()
        transient_federation_attribute = utils.get_setting(idp, 'TRANSIENT_FEDERATION_ATTRIBUTE')
        if saml_attributes['name_id_format'] == lasso.SAML2_NAME_IDENTIFIER_FORMAT_TRANSIENT:
            if (transient_federation_attribute
                    and saml_attributes.get(transient_federation_attribute)):
                name_id = saml_attributes[transient_federation_attribute]
                if not isinstance(name_id, basestring):
                    if len(name_id) == 1:
                        name_id = name_id[0]
                    else:
                        self.logger.warning('more than one value for attribute %r, cannot federate',
                                            transient_federation_attribute)
                        return None
            else:
                return None
        else:
            name_id = saml_attributes['name_id_content']
        issuer = saml_attributes['issuer']
        try:
            return User.objects.get(saml_identifiers__name_id=name_id,
                                    saml_identifiers__issuer=issuer)
        except User.DoesNotExist:
            if not utils.get_setting(idp, 'PROVISION'):
                self.logger.warning('provisionning disabled, login refused')
                return None
            user = self.create_user(User)
            saml_id, created = models.UserSAMLIdentifier.objects.get_or_create(
                name_id=name_id, issuer=issuer, defaults={'user': user})
            if created:
                try:
                    self.finish_create_user(idp, saml_attributes, user)
                except UserCreationError:
                    user.delete()
                    return None
                self.logger.info('created new user %s with name_id %s from issuer %s',
                                 user, name_id, issuer)
            else:
                user.delete()
                user = saml_id.user
                self.logger.info('looked up user %s with name_id %s from issuer %s',
                                 user, name_id, issuer)
        return user

    def provision(self, user, idp, saml_attributes):
        self.provision_attribute(user, idp, saml_attributes)
        self.provision_superuser(user, idp, saml_attributes)
        self.provision_groups(user, idp, saml_attributes)

    def provision_attribute(self, user, idp, saml_attributes):
        realm = utils.get_setting(idp, 'REALM')
        attribute_mapping = utils.get_setting(idp, 'ATTRIBUTE_MAPPING')
        attribute_set = False
        for field, tpl in attribute_mapping.iteritems():
            try:
                value = unicode(tpl).format(realm=realm, attributes=saml_attributes, idp=idp)
            except ValueError:
                self.logger.warning(u'invalid attribute mapping template %r', tpl)
            except (AttributeError, KeyError, IndexError, ValueError) as e:
                self.logger.warning(
                    u'invalid reference in attribute mapping template %r: %s', tpl, e)
            else:
                model_field = user._meta.get_field(field)
                if hasattr(model_field, 'max_length'):
                    value = value[:model_field.max_length]
                if getattr(user, field) != value:
                    old_value = getattr(user, field)
                    setattr(user, field, value)
                    attribute_set = True
                    self.logger.info(u'set field %s of user %s to value %r (old value %r)', field,
                                     user, value, old_value)
        if attribute_set:
            user.save()

    def provision_superuser(self, user, idp, saml_attributes):
        superuser_mapping = utils.get_setting(idp, 'SUPERUSER_MAPPING')
        if not superuser_mapping:
            return
        attribute_set = False
        for key, values in superuser_mapping.iteritems():
            if key in saml_attributes:
                if not isinstance(values, (tuple, list)):
                    values = [values]
                values = set(values)
                attribute_values = saml_attributes[key]
                if not isinstance(attribute_values, (tuple, list)):
                    attribute_values = [attribute_values]
                attribute_values = set(attribute_values)
                if attribute_values & values:
                    if not (user.is_staff and user.is_superuser):
                        user.is_staff = True
                        user.is_superuser = True
                        attribute_set = True
                        self.logger.info('flag is_staff and is_superuser added to user %s', user)
                    break
        else:
            if user.is_superuser or user.is_staff:
                user.is_staff = False
                user.is_superuser = False
                self.logger.info('flag is_staff and is_superuser removed from user %s', user)
                attribute_set = True
        if attribute_set:
            user.save()

    def provision_groups(self, user, idp, saml_attributes):
        User = user.__class__
        group_attribute = utils.get_setting(idp, 'GROUP_ATTRIBUTE')
        create_group = utils.get_setting(idp, 'CREATE_GROUP')
        if group_attribute in saml_attributes:
            values = saml_attributes[group_attribute]
            if not isinstance(values, (list, tuple)):
                values = [values]
            groups = []
            for value in set(values):
                if create_group:
                    group, created = Group.objects.get_or_create(name=value)
                else:
                    try:
                        group = Group.objects.get(name=value)
                    except Group.DoesNotExist:
                        continue
                groups.append(group)
            for group in Group.objects.filter(pk__in=[g.pk for g in groups]).exclude(user=user):
                self.logger.info(
                    u'adding group %s (%s) to user %s (%s)', group, group.pk, user, user.pk)
                User.groups.through.objects.get_or_create(group=group, user=user)
            qs = User.groups.through.objects.exclude(
                group__pk__in=[g.pk for g in groups]).filter(user=user)
            for rel in qs:
                self.logger.info(u'removing group %s (%s) from user %s (%s)', rel.group,
                                 rel.group.pk, rel.user, rel.user.pk)
            qs.delete()