import os
import sys
import urlparse

try:
    import lasso
except ImportError:
    print >> sys.stderr, 'Missing Lasso module, SAMLv2 support disabled'

from quixote import get_publisher, get_request, get_response, get_session, get_session_manager, redirect

from qommon.liberty import SOAPException, soap_call
from qommon.saml2 import Saml2Directory
from qommon import template
from qommon import get_logger

import misc
from users import User
from hosts import Host
from federations import Federation
import site_authentication

class Saml2(Saml2Directory):
    _q_exports = Saml2Directory._q_exports + ['local_auth']

    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()
        #remote_provider_cfg = get_cfg('idp', {}).get(misc.get_provider_key(login.remoteProviderId))
        #client_cert = remote_provider_cfg.get('clientcertificate')

        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):
        providerId = login.server.providerId
        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 != 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_dumps[providerId] = login.session.dump()
                session.lasso_session_indexes[providerId] = assertion.authnStatement[0].sessionIndex
                session.lasso_session_name_identifiers[providerId] = login.nameIdentifier.content
            else:
                session.lasso_session_dumps[login.server.providerId] = None

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

        user = self.lookup_user(session, login)

        # Check if it is for Larpe administration or token
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])
        if host.name == 'larpe':
            if user:
                session.set_user(user.id, login.server.providerId)
            else:
                session.name_identifier = login.nameIdentifier.content
                session.lasso_anonymous_identity_dump = login.identity.dump()
                session.provider_id = login.server.providerId

            if session.after_url:
                # Access to an admin page or token url with parameter
                after_url = session.after_url
                session.after_url = None
                return redirect(after_url)

            if user and user.is_admin:
                return redirect('%s/admin/' % get_request().environ['SCRIPT_NAME'])
            else:
                return redirect('%s/token' % get_request().environ['SCRIPT_NAME'])

        # Set session user
        if not user:
            user = User()
        user.name_identifiers = [ login.nameIdentifier.content ]
        user.lasso_dumps = [ login.identity.dump() ]
        user.store()
        session.set_user(user.id, login.server.providerId)

        # Check if a federation already exist
        federations = Federation.select(lambda x: host.id == x.host_id \
                                        and user.name_identifiers[0] in x.name_identifiers)

        if federations:
            return site_authentication.get_site_authentication(host).sso_local_login(federations[0])
        else:
            # Build response redirection
            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('local_auth'))
            response.content_type = 'text/plain'
            return 'Your browser should redirect you'

    def lookup_user(self, session, login):
        found_users = list(User.select(lambda x: login.nameIdentifier.content in x.name_identifiers, ignore_errors = True))
        if found_users:
            return found_users[0]
        return None

    def slo_sp(self, method = None):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])

        if method is None:
            method = lasso.HTTP_METHOD_REDIRECT

        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        session = get_session()

        if not session.id or not session.users.has_key(logout.server.providerId) \
                or not session.lasso_session_dumps.has_key(logout.server.providerId):
            get_session_manager().expire_session(logout.server.providerId)
            return redirect(host.get_root_url())
        logout.setSessionFromDump(session.lasso_session_dumps[logout.server.providerId])
        user = session.get_user(logout.server.providerId)

        if host.name != 'larpe' and user:
            site_authentication.get_site_authentication(host).local_logout(user=user)

        if user and user.lasso_dumps:
            logout.setIdentityFromDump(user.lasso_dumps[0])
        else:
            get_session_manager().expire_session(logout.server.providerId)
            return redirect(host.get_root_url())

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

        # Not implemented yet
        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(logout.server.providerId)
                return redirect(host.get_root_url())
            if error[0] == lasso.PROFILE_ERROR_SESSION_NOT_FOUND:
                get_session_manager().expire_session(logout.server.providerId)
                return redirect(host.get_root_url())
            raise

        logout.buildRequestMsg()
        return redirect(logout.msgUrl)

    def singleLogoutReturn(self):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])

        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        session = get_session()

        if not session.id or not session.users.has_key(logout.server.providerId) \
                or not session.lasso_session_dumps.has_key(logout.server.providerId):
            get_session_manager().expire_session(logout.server.providerId)
            return redirect(host.get_root_url())
        logout.setSessionFromDump(session.lasso_session_dumps[logout.server.providerId])

        message = get_request().get_query()
        return self.slo_return(logout, message)

    def slo_return(self, logout, message):
        host = Host.get_host_from_url()

        session = get_session()

        try:
            logout.processResponseMsg(message)
        except lasso.Error, error:
            if error[0] == lasso.PROFILE_ERROR_INVALID_QUERY:
                get_logger().warn('Invalid response')
            elif error[0] == lasso.DS_ERROR_INVALID_SIGNATURE:
                get_logger().warn('Failed to check single logout request signature')
            elif error[0] == lasso.LOGOUT_ERROR_REQUEST_DENIED:
                get_logger().warn('Request Denied')
            elif error[0] == lasso.LOGOUT_ERROR_UNKNOWN_PRINCIPAL:
                get_logger().warn('Unknown principal on logout, probably session stopped already on IdP')
            else:
                get_logger().error('Unknown Lasso exception on logout return: ' + repr(error))
        except Exception, exception:
            get_logger().error('Unknown exception on logout return: ' + repr(exception))

        get_session_manager().expire_session(logout.server.providerId)

        return redirect(host.get_root_url())

    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:
            get_logger().warn('SOAP message on single logout url not a slo message')
            return

        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        providerId = logout.server.providerId
        logout.processRequestMsg(soap_message)
        name_identifier = logout.nameIdentifier.content
        # find one session matching the name identifier, and eventually the request
        for session in get_session_manager().values():
            session_index = session.lasso_session_indexes.get(providerId)
            name_identifier = session.lasso_session_name_identifiers.get(providerId)
            request_name_identifier = logout.nameIdentifier.content
            request_session_index = logout.request.sessionIndex
            if request_name_identifier == name_identifier and \
                    (not session_index or request_session_index == session_index) \
                        and session.lasso_session_dumps.get(providerId):
                get_logger().info('SLO/SOAP from %s' % logout.remoteProviderId)
                break
        else:
            # no session, build straight failure answer
            logout.buildResponseMsg()
            return logout.msgBody

        return self.slo_idp(logout, session)

    def singleLogout(self):
        logout = lasso.Logout(misc.get_lasso_server(protocol = 'saml2'))
        try:
            logout.processRequestMsg(get_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
            # find a matching 
            for session in get_session_manager().values():
                session_index = session.lasso_session_indexes.get(providerId)
                name_identifier = session.lasso_session_name_identifiers.get(providerId)
                request_name_identifier = logout.nameIdentifier.content
                request_session_index = logout.request.sessionIndex
                if request_name_identifier == name_identifier and \
                        (not session_index or request_session_index == session_index) \
                        and session.lasso_session_dump.get(providerId):
                    get_logger().info('SLO/SOAP from %s' % logout.remoteProviderId)
                    break
            else:
                # no session, build straight failure answer
                logout.buildResponseMsg()
                return logout.msgBody

        return self.slo_idp(logout, session)

    def slo_idp(self, logout, session):
        # This block differs from qommon
        if session.lasso_session_dumps.has_key(logout.server.providerId):
            logout.setSessionFromDump(session.lasso_session_dumps[logout.server.providerId])
        user = session.get_user(logout.server.providerId)
        if user and user.lasso_dumps:
            logout.setIdentityFromDump(user.lasso_dumps[0])

        if user and logout.nameIdentifier.content not in user.name_identifiers:
            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
            elif error[0] == lasso.SERVER_ERROR_PROVIDER_NOT_FOUND:
                pass
            elif error[0] == lasso.NAME_IDENTIFIER_NOT_FOUND:
                pass
            else:
                raise
        else:
            providerId = logout.server.providerId
            session_index = logout.request.sessionIndex
            name_identifier = logout.nameIdentifier.content
            # Remove reference to local authentication on this SP in the session
            # if a user is present, try a local logout
            for session2 in get_session_manager().values():
                if session2.lasso_session_name_identifiers.get(providerId) == name_identifier \
                        and ( not session_index
                                or session2.lasso_session_indexes.get(providerId) == session_index):
                    if session2.users.has_key(providerId):
                        # local logout
                        site_auth = site_authentication.get_site_authentication(Host.get_host_from_url())
                        site_auth.local_logout(user=session2.get_user(providerId),
                                cookies=getattr(session2,'cookies', None))
                        del session2.users[providerId]
                    if session2.lasso_session_dumps.has_key(providerId):
                        del session2.lasso_session_dumps[providerId]
                    if session2.lasso_session_indexes.has_key(providerId):
                        del session2.lasso_session_indexes[providerId]
                    if session2.lasso_session_name_identifiers.has_key(providerId):
                        del session2.lasso_session_name_identifiers[providerId]
                    session2.store()
            get_session_manager().expire_session(logout.server.providerId)

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

    def local_auth(self):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])
        return site_authentication.get_site_authentication(host).local_auth
    local_auth = property(local_auth)

    def metadata(self):
        host = Host.get_host_from_url()
        if host is None:
            return redirect('%s/' % get_request().environ['SCRIPT_NAME'])
        get_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'])
        get_response().set_content_type('text/plain')
        public_key = open(host.public_key).read()
        return public_key

