package com.onaro.sanscreen.client.login;

import java.io.IOException;
import java.security.Principal;
import java.util.Map;
import java.util.Set;

import javax.security.auth.Subject;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;

import org.jboss.security.SecurityContext;
import org.jboss.security.SecurityContextAssociation;
import org.jboss.security.SecurityContextFactory;
import org.jboss.security.SimplePrincipal;
import org.jboss.security.auth.callback.ObjectCallback;

/**
 * This login module does much of what the ClientLoginClass does, but uses the CallbackHandler to retrieve a certificate
 * which it then uses as the credential when setting the principalInfo of the onaro-login SecurityContext.<p/>
 *
 * @author Joshua Eisenthal  Copyright 2014 NetApp
 */
public class ClientCertLoginModule implements LoginModule {
    private Subject subject;
    private CallbackHandler callbackHandler;
    private String username;
    private Object certificate;
    private SimplePrincipal loginPrincipal;

    @Override
    public void initialize(Subject aSubject, CallbackHandler aCallbackHandler, Map<String, ?> sharedState, Map<String, ?> options) {
        subject = aSubject;
        callbackHandler = aCallbackHandler;
    }

    @Override
    public boolean login() throws LoginException {
        if (callbackHandler == null) {
            throw new LoginException("CallbackHandler not set");     //$NON-NLS-1$
        }

        NameCallback nameCallback = new NameCallback("Username: "); //$NON-NLS-1$
        ObjectCallback certCallback = new ObjectCallback("Certificate: "); //$NON-NLS-1$
        Callback[] callbacks = {nameCallback, certCallback};
        try {
            callbackHandler.handle(callbacks);
        } catch (IOException e) {
            LoginException le = new LoginException(e.getMessage());
            le.initCause(e);
            throw le;
        } catch (UnsupportedCallbackException e) {
            LoginException le = new LoginException("CallbackHandler does not support username/certificate callbacks"); //$NON-NLS-1$
            le.initCause(e);
            throw le;
        }
        username = nameCallback.getName();
        certificate = certCallback.getCredential();

        return true;
    }

    @Override
    public boolean commit() throws LoginException {
        loginPrincipal = new SimplePrincipal(username);
        try {
            // This creates the client security context. It is somehow remotely connected to the server's
            // "onaro-login" realm, and thus by specifying the principal and credential, invokes the server's
            // login modules:
            SecurityContext securityContext = SecurityContextFactory.createSecurityContext(loginPrincipal, certificate, subject, "CLIENT_LOGIN_MODULE"); //$NON-NLS-1$
            SecurityContextAssociation.setSecurityContext(securityContext);
        } catch (Exception e) {
            LoginException le = new LoginException("Failed to get security context"); //$NON-NLS-1$
            le.initCause(e);
            throw le;
        }

        Set<Principal> principals = subject.getPrincipals();
        if (!principals.contains(loginPrincipal)) {
            principals.add(loginPrincipal);
        }

        return true;
    }

    @Override
    public boolean abort() throws LoginException {
        return logout();
    }

    @Override
    public boolean logout() throws LoginException {
        SecurityContextAssociation.clearSecurityContext();

        Set<Principal> principals = subject.getPrincipals();
        principals.remove(loginPrincipal);

        return true;
    }
}