package cz.muni.ics.oidc.saml;

import lombok.extern.slf4j.Slf4j;
import org.joda.time.DateTime;
import org.opensaml.saml2.core.AuthnContext;
import org.opensaml.saml2.core.AuthnContextClassRef;
import org.opensaml.saml2.core.AuthnStatement;
import org.springframework.security.core.Authentication;
import org.springframework.security.providers.ExpiringUsernameAuthenticationToken;
import org.springframework.security.saml.SAMLCredential;
import org.springframework.security.web.authentication.SavedRequestAwareAuthenticationSuccessHandler;
import org.springframework.security.web.authentication.WebAuthenticationDetails;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.stream.Collectors;

@Slf4j
public class PerunSamlAuthenticationSuccessHandler extends SavedRequestAwareAuthenticationSuccessHandler {

    public static final String SESSION_ACR = "ACR";

    public static final String SESSION_AUTH_TIMESTAMP = "AUTH_TIMESTAMP";

    @Override
    public void onAuthenticationSuccess(HttpServletRequest request,
                                        HttpServletResponse response,
                                        Authentication authentication)
            throws ServletException, IOException
    {
        if (authentication instanceof ExpiringUsernameAuthenticationToken) {
            ExpiringUsernameAuthenticationToken token = (ExpiringUsernameAuthenticationToken) authentication;
            Object details = token.getDetails();
            if (details instanceof WebAuthenticationDetails) {
                WebAuthenticationDetails webDetails = (WebAuthenticationDetails) details;
                log.info("successful authentication, remote IP address {}", webDetails.getRemoteAddress());
            }
            if (token.getCredentials() instanceof SAMLCredential) {
                SAMLCredential credential = (SAMLCredential) token.getCredentials();
                setAcrToSession(request, credential);
                setAuthTime(request, credential);
            }
        }
        super.onAuthenticationSuccess(request, response, authentication);
    }

    private void setAcrToSession(HttpServletRequest request,  SAMLCredential credential) {
        String acrs = credential.getAuthenticationAssertion()
                .getAuthnStatements().stream()
                .map(AuthnStatement::getAuthnContext)
                .map(AuthnContext::getAuthnContextClassRef)
                .map(AuthnContextClassRef::getAuthnContextClassRef)
                .collect(Collectors.joining());
        request.getSession(true).setAttribute(SESSION_ACR, acrs);
    }

    private void setAuthTime(HttpServletRequest req, SAMLCredential credential) {
        DateTime authnInstant = credential.getAuthenticationAssertion().getAuthnStatements().get(0).getAuthnInstant();
        req.getSession(true).setAttribute(SESSION_AUTH_TIMESTAMP, authnInstant.getMillis());
    }

}
