import os
import urlparse
import time

import lasso

from quixote import get_request, get_response, redirect, get_field, get_publisher
from quixote.http_request import parse_header
from quixote.directory import Directory
from quixote import get_user, get_session, get_session_manager

import logger
import misc
import template

from liberty import SOAPException, soap_call


class Saml2Directory(Directory):
    _q_exports = ['login',
            'singleSignOnArtifact', 'singleSignOnPost', 'singleSignOnSOAP',
            'singleLogout', 'singleLogoutReturn', 'singleLogoutSOAP',
            'manageNameId', 'manageNameIdReturn', 'manageNameIdSOAP',
            'metadata', ('metadata.xml', 'metadata'), 'public_key']

    def _q_index(self):
        raise errors.AccessError()

    def singleSignOnSOAP(self):
        request = get_request()
        server = misc.get_lasso_server(protocol = 'saml2')
        if not server:
            return template.error_page(_('SAML 2.0 support not yet configured.'))
        login = lasso.Login(server)
        paos_message = self.get_soap_message()
        try:
            login.processPaosResponseMsg(paos_message)
        except lasso.Error, error:
            if error[0] == lasso.LOGIN_ERROR_STATUS_NOT_SUCCESS:
                response = get_response()
                response.set_status(401)
                return 'Unknown authentication failure'
            raise
        login.acceptSso()  

        return self.success_ecp()

    def success_ecp [html] (self):
        template.html_top()
        'ECP Authentication succeeded.'

    def login(self):
        return self.perform_login()

    def perform_login(self, idp = None):
        server = misc.get_lasso_server(protocol = 'saml2')
        if not server:
            return template.error_page(_('SAML 2.0 support not yet configured.'))
        login = lasso.Login(server)
        login.initAuthnRequest(idp, lasso.HTTP_METHOD_REDIRECT)
        login.request.nameIDPolicy.format = lasso.SAML2_NAME_IDENTIFIER_FORMAT_PERSISTENT
        login.request.nameIDPolicy.allowCreate = True
        login.request.forceAuthn = False
        login.request.isPassive = False
        login.request.consent = 'urn:oasis:names:tc:SAML:2.0:consent:current-implicit'
        login.buildAuthnRequestMsg()
        return redirect(login.msgUrl)
        

    def singleSignOnArtifact(self):
        server = misc.get_lasso_server(protocol = 'saml2')
        if not server:
            return template.error_page(_('SAML 2.0 support not yet configured.'))
        login = lasso.Login(server)
        request = get_request()
        try:
            login.initRequest(request.get_query(), lasso.HTTP_METHOD_ARTIFACT_GET)
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_MISSING_ARTIFACT:
                return template.error_page(_('Missing SAML Artifact'))
            else:
                raise

        login.buildRequestMsg()
        try:
            soap_answer = soap_call(login.msgUrl, login.msgBody)
        except SOAPException:
            return template.error_page(_('Failure to communicate with identity provider'))

        try:
            login.processResponseMsg(soap_answer)
        except lasso.Error, error:
            if error[0] == lasso.LOGIN_ERROR_STATUS_NOT_SUCCESS:
                return template.error_page(_('Unknown authentication failure'))
            if error[0] == lasso.LOGIN_ERROR_UNKNOWN_PRINCIPAL:
                return template.error_page(_('Authentication failure; unknown principal'))
            if error[0] == lasso.LOGIN_ERROR_FEDERATION_NOT_FOUND:
                return template.error_page('there was no federation')
            raise

        return self.sso_after_response(login)

    def sso_after_response(self, login):
        try:
            assertion = login.response.assertion[0]
            if assertion.subject.subjectConfirmation.subjectConfirmationData.recipient != \
                        get_request().get_url():
                return template.error_page('SubjectConfirmation Recipient Mismatch')
        except:
                return template.error_page('SubjectConfirmation Recipient Mismatch')

        assertions_dir = os.path.join(get_publisher().app_dir, 'assertions')
        if not os.path.exists(assertions_dir):
            os.mkdir(assertions_dir)

        assertion_fn = os.path.join(assertions_dir, assertion.iD)
        if os.path.exists(assertion_fn):
            return template.error_page('Assertion replay')

        try:
            if assertion.subject.subjectConfirmation.method != \
                        'urn:oasis:names:tc:SAML:2.0:cm:bearer':
                return template.error_page('Unknown SubjectConfirmation Method')
        except:
            return template.error_page('Unknown SubjectConfirmation Method')

        try:
            audience_ok = False
            for audience_restriction in assertion.conditions.audienceRestriction:
                if audience_restriction.audience != login.server.providerId:
                    return template.error_page('Incorrect AudienceRestriction')
                audience_ok = True
            if not audience_ok:
                return template.error_page('Incorrect AudienceRestriction')
        except:
            return template.error_page('Incorrect AudienceRestriction')

        try:
            current_time = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
            not_before = assertion.subject.subjectConfirmation.subjectConfirmationData.notBefore
            not_on_or_after = assertion.subject.subjectConfirmation.subjectConfirmationData.notOnOrAfter
            if not_before and current_time < not_before:
                return template.error_page('Assertion received too early')
            if not_on_or_after and current_time > not_on_or_after:
                return template.error_page('Assertion expired')
        except:
            return template.error_page('Error checking Assertion Time')

        # TODO: check for unknown conditions

        login.acceptSso()
        session = get_session()
        if login.isSessionDirty:
            if login.session:
                session.lasso_session_dump = login.session.dump()
            else:
                session.lasso_session_dump = None

        if assertion.authnStatement[0].sessionIndex:
            session.lasso_session_index = assertion.authnStatement[0].sessionIndex

        user = self.lookup_user(session, login)
        if user:
            session.set_user(user.id)
        else:
            session.set_user('anonymous-%s' % login.nameIdentifier.content)
            if login.identity:
                session.lasso_anonymous_identity_dump = login.identity.dump()
            else:
                # XXX: this situation happened with SSO initiated by IdP, this
                # is not normal
                pass
        session.lasso_identity_provider_id = login.remoteProviderId

        response = get_response()
        if session.after_url:
            after_url = session.after_url
            session.after_url = None
            return redirect(after_url)
        response.set_status(303)
        response.headers['location'] = urlparse.urljoin(get_request().get_url(), str('..'))
        response.content_type = 'text/plain'
        return "Your browser should redirect you"


    def singleSignOnPost(self):
        server = misc.get_lasso_server(protocol = 'saml2')
        if not server:
            return template.error_page(_('SAML 2.0 support not yet configured.'))
        login = lasso.Login(server)

        if not get_field('SAMLResponse'):
            return template.error_page(_('No SAML Response'))

        try:
            login.processAuthnResponseMsg(get_field('SAMLResponse'))
        except lasso.Error, error:
            if error[0] == lasso.LOGIN_ERROR_STATUS_NOT_SUCCESS:
                return template.error_page(_('Unknown authentication failure'))
            if error[0] == lasso.SERVER_ERROR_PROVIDER_NOT_FOUND:
                return template.error_page(_('Request from unknown provider ID'))
            if error[0] == lasso.DS_ERROR_SIGNATURE_NOT_FOUND:
                return template.error_page(_('Error checking signature'))
            if error[0] == lasso.DS_ERROR_SIGNATURE_VERIFICATION_FAILED:
                return template.error_page(_('Error checking signature'))
            if error[0] == lasso.DS_ERROR_INVALID_SIGNATURE:
                return template.error_page(_('Error checking signature'))
            raise

        return self.sso_after_response(login)


    def lookup_user(self, session, login):
        ni = login.nameIdentifier.content
        session.name_identifier = ni
        nis = list(get_publisher().user_class.select(lambda x: ni in x.name_identifiers))
        if nis:
            user = nis[0]
        else:
            return None

        user.lasso_dump = login.identity.dump()
        user.store()
        return user

    def slo_sp(self, method = lasso.HTTP_METHOD_REDIRECT):
        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        session = get_session()

        if session.lasso_session_dump:
            logout.setSessionFromDump(session.lasso_session_dump)

        user = get_request().user
        if user and user.lasso_dump:
            logout.setIdentityFromDump(user.lasso_dump)

        if method == lasso.HTTP_METHOD_REDIRECT:
            return self.slo_sp_redirect(logout)

        if method == lasso.HTTP_METHOD_SOAP:
            return self.slo_sp_soap(logout)

    def slo_sp_redirect(self, logout):
        session = get_session()
        try:
            logout.initRequest(None, lasso.HTTP_METHOD_REDIRECT)
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_NAME_IDENTIFIER_NOT_FOUND:
                get_session_manager().expire_session() 
                return redirect(get_request().environ['SCRIPT_NAME'] + '/')
            if error[0] == lasso.PROFILE_ERROR_SESSION_NOT_FOUND:
                get_session_manager().expire_session() 
                return redirect(get_request().environ['SCRIPT_NAME'] + '/')
            raise
        if session.lasso_session_index:
            logout.request.sessionIndex = session.lasso_session_index
        logout.buildRequestMsg()
        return redirect(logout.msgUrl)


    def singleLogoutReturn(self):
        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        if get_session().lasso_session_dump:
            logout.setSessionFromDump(get_session().lasso_session_dump)
        message = get_request().get_query()
        return self.slo_return(logout, message)

    def slo_return(self, logout, message):
        session = get_session()

        try:
            logout.processResponseMsg(message)
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_INVALID_QUERY:
                return template.error_page(_('Invalid response'))
            elif error[0] == lasso.DS_ERROR_INVALID_SIGNATURE:
                return template.error_page(_('Failed to check single logout request signature.'))
            elif error[0] == lasso.LOGOUT_ERROR_REQUEST_DENIED:
                return template.error_page(_('Request Denied'))
            elif error[0] == lasso.LOGOUT_ERROR_UNKNOWN_PRINCIPAL:
                logger.warn('Unknown principal on logout, probably session stopped already on IdP')
                # XXX: wouldn't work when logged on two IdP
                session.lasso_session_dump = None
            else:
                raise
            
        if logout.isSessionDirty:
            if logout.session:
                session.lasso_session_dump = logout.session.dump()
            else:
                session.lasso_session_dump = None

        if not session.lasso_session_dump:
            get_session_manager().expire_session() 

        return redirect(get_request().environ['SCRIPT_NAME'] + '/')


    def slo_sp_soap(self, logout):
        try:
            logout.initRequest(None, lasso.HTTP_METHOD_SOAP)
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_NAME_IDENTIFIER_NOT_FOUND:
                get_session_manager().expire_session() 
                return redirect(get_request().environ['SCRIPT_NAME'] + '/')
            if error[0] == lasso.PROFILE_ERROR_SESSION_NOT_FOUND:
                get_session_manager().expire_session() 
                return redirect(get_request().environ['SCRIPT_NAME'] + '/')
            raise
        logout.buildRequestMsg()
        try:
            soap_answer = soap_call(logout.msgUrl, logout.msgBody)
        except SOAPException:
            return template.error_page(_('Failure to communicate with identity provider'))

        return self.slo_return(logout, soap_answer)


    def fedterm_sp(self, method = lasso.HTTP_METHOD_REDIRECT):
        manage = lasso.NameIdManagement(misc.get_lasso_server(protocol = 'saml2'))
        session = get_session()

        if session.lasso_session_dump:
            manage.setSessionFromDump(session.lasso_session_dump)

        user = get_request().user
        if user and user.lasso_dump:
            manage.setIdentityFromDump(user.lasso_dump)

        remote_provider_id = get_session().lasso_identity_provider_id

        if method == lasso.HTTP_METHOD_REDIRECT:
            return self.fedterm_sp_redirect(manage, remote_provider_id)

        if method == lasso.HTTP_METHOD_SOAP:
            return self.fedterm_sp_soap(manage, remote_provider_id)

    def get_soap_message(self):
        request = get_request()
        ctype = request.environ.get('CONTENT_TYPE')
        if not ctype:
            logger.warn('SOAP Endpoint got message without content-type')
            raise SOAPError()

        ctype, ctype_params = parse_header(ctype)
        if ctype not in ('text/xml', 'application/vnd.paos+xml'):
            logger.warn('SOAP Endpoint got message with wrong content-type (%s)' % ctype)
            raise SOAPError()

        length = int(request.environ.get('CONTENT_LENGTH'))
        return request.stdin.read(length)


    def singleLogoutSOAP(self):
        try:
            soap_message = self.get_soap_message()
        except:
            return

        response = get_response()
        response.set_content_type('text/xml')

        request_type = lasso.getRequestTypeFromSoapMsg(soap_message) 

        if request_type != lasso.REQUEST_TYPE_LOGOUT:
            logger.warn('SOAP message on single logout url not a slo message')
            return

        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        logout.processRequestMsg(soap_message)
        name_identifier = logout.nameIdentifier.content
        for session in get_session_manager().values():
            if session.name_identifier == name_identifier:
                logger.info('SLO/SOAP from %s' % logout.remoteProviderId, session = session)
                logout.setSessionFromDump(session.lasso_session_dump)
                break
        else:
            # no session, build straight failure answer
            logout.buildResponseMsg()
            return logout.msgBody

        user = session.get_user()
        if user and user.lasso_dump:
            logout.setIdentityFromDump(user.lasso_dump)

        return self.slo_idp(logout, session)

    def singleLogout(self):
        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        request = get_request()
        try:
            logout.processRequestMsg(request.get_query())
        except lasso.Error, error:
            if error[0] == lasso.DS_ERROR_INVALID_SIGNATURE:
                return template.error_page(_('Failed to check single logout request signature.'))
            raise
        session = get_session()
        if not session.id:
            # session has not been found, this may be because the user has
            # its browser configured so that cookies are not sent for
            # remote queries and IdP is using image-based SLO.
            # so we look up a session with the appropriate name identifier
            name_identifier = logout.nameIdentifier.content
            for session in get_session_manager().values():
                if name_identifier == session.name_identifier:
                    break
            else:
                session = get_session()

        return self.slo_idp(logout, session)

    def slo_idp(self, logout, session):
        if session.lasso_session_dump:
            logout.setSessionFromDump(session.lasso_session_dump)
        user = session.get_user()
        if user and user.lasso_dump:
            logout.setIdentityFromDump(user.lasso_dump)
        if logout.nameIdentifier.content != session.name_identifier:
            raise 'no appropriate name identifier in session (%s and %s)' % (
                    logout.nameIdentifier.content, session.name_identifier)

        try:
            assertion = logout.session.getAssertions(logout.remoteProviderId)[0]
            if logout.request.sessionIndex and (
                    assertion.authnStatement[0].sessionIndex != logout.request.sessionIndex):
                logout.setSessionFromDump('<Session />')
        except:
            pass
        
        try:
            logout.validateRequest()
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_SESSION_NOT_FOUND:
                pass
            elif error[0] == lasso.PROFILE_ERROR_IDENTITY_NOT_FOUND:
                pass
            elif error[0] == lasso.PROFILE_ERROR_MISSING_ASSERTION:
                pass
            else:
                raise
        else:
            get_session_manager().expire_session() 
            try:
                del get_session_manager()[session.id]
            except KeyError:
                pass

        try:
            if not logout.request.sessionIndex:
                for session2 in get_session_manager().values():
                    if name_identifier == session2.name_identifier:
                        del get_session_manager()[session2.id]
        except:
            # killing all session failed, ignoring silently
            pass

        logout.buildResponseMsg()
        if logout.msgBody: # soap answer
            return logout.msgBody
        else:
            return redirect(logout.msgUrl)





    def fedterm_sp_redirect(self, manage, remote_provider_id):
        manage.initRequest(remote_provider_id, None, lasso.HTTP_METHOD_REDIRECT)
        manage.buildRequestMsg()
        get_session().lasso_manage_name_id_dump = manage.dump()
        return redirect(manage.msgUrl)

    def fedterm_sp_soap(self, manage, remote_provider_id):
        manage.initRequest(remote_provider_id, None, lasso.HTTP_METHOD_SOAP)
        manage.buildRequestMsg()
        try:
            soap_answer = soap_call(manage.msgUrl, manage.msgBody)
        except SOAPException:
            return template.error_page(_('Failure to communicate with identity provider'))

        return self.manage_name_id_return(manage, soap_answer)

    def manageNameId(self):
        manage = lasso.NameIdManagement(misc.get_lasso_server(protocol = 'saml2'))
        manage.processRequestMsg(get_request().get_query())

        session = get_session()
        user = get_request().user
        self.manage_name_id(manage, user, session)
        return redirect(manage.msgUrl)

    def manageNameIdSOAP(self):
        try:
            soap_message = self.get_soap_message()
        except:
            return

        response = get_response()
        response.set_content_type('text/xml')

        request_type = lasso.getRequestTypeFromSoapMsg(soap_message) 

        if request_type != lasso.REQUEST_TYPE_NAME_ID_MANAGEMENT:
            logger.warn('SOAP message on name id management url of not appropriate type')
            return

        manage = lasso.NameIdManagement(misc.get_lasso_server(protocol = 'saml2'))
        manage.processRequestMsg(soap_message)

        ni = manage.nameIdentifier.content
        nis = list(get_publisher().user_class.select(lambda x: ni in x.name_identifiers))
        session = None
        if nis:
            user = nis[0]
        else:
            user = None
            for session in get_session_manager().values():
                if session.name_identifier == ni:
                    user = session.get_user()
                    break

        self.manage_name_id(manage, user, session)
        return manage.msgBody


    def manage_name_id(self, manage, user, session):
        if user and user.lasso_dump:
            manage.setIdentityFromDump(user.lasso_dump)
        elif session and session.lasso_anonymous_identity_dump:
            manage.setIdentityFromDump(session.lasso_anonymous_identity_dump)

        try:
            manage.validateRequest()
        except lasso.Error, error:
            raise
        else:
            if not manage.identity:
                user.lasso_dump = None
            else:
                user.lasso_dump = manage.identity.dump()
            if user.anonymous:
                session.lasso_anonymous_identity_dump = user.lasso_dump
                get_session_manager().maintain_session(session)
            else:
                user.store()

        manage.buildResponseMsg()

    def manageNameIdReturn(self):
        if get_session().lasso_manage_name_id_dump:
            manage = lasso.NameIdManagement.newFromDump(
                    misc.get_lasso_server(protocol = 'saml2'),
                    get_session().lasso_manage_name_id_dump)
            get_session().lasso_manage_name_id_dump = None
        else:
            manage = lasso.NameIdManagement(misc.get_lasso_server(protocol = 'saml2'))

        user = get_request().user
        if user and user.lasso_dump:
            manage.setIdentityFromDump(user.lasso_dump)

        message = get_request().get_query()
        return self.manage_name_id_return(manage, message)

    def manage_name_id_return(self, manage, message):
        try:
            manage.processResponseMsg(message)
        except lasso.Error, error:
            raise

        if manage.isIdentityDirty:
            if manage.identity:
                get_request().user.lasso_dump = manage.identity.dump()
            else:
                get_request().user.lasso_dump = None

        if not get_request().user.lasso_dump:
            get_session_manager().expire_session() 

        return redirect(get_request().environ['SCRIPT_NAME'] + '/')

    def metadata(self):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])
        response = get_response()
        response.set_content_type('text/xml', 'utf-8')
        metadata = unicode(open(host.saml2_metadata).read(), 'utf-8')
        return metadata

    def public_key(self):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])
        response = get_response()
        response.set_content_type('application/octet-stream')
        public_key = open(host.public_key).read()
        return public_key
