/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.wss4j.common.spnego;

import java.security.Principal;
import java.security.PrivilegedActionException;
import java.util.Set;

import javax.security.auth.Subject;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.login.LoginContext;
import javax.security.auth.login.LoginException;

import org.apache.wss4j.common.ext.WSSecurityException;
import org.apache.wss4j.common.ext.WSSecurityException.ErrorCode;
import org.apache.wss4j.common.kerberos.KerberosClientExceptionAction;
import org.apache.wss4j.common.kerberos.KerberosContext;
import org.apache.wss4j.common.kerberos.KerberosServiceContext;
import org.apache.wss4j.common.kerberos.KerberosServiceExceptionAction;
import org.ietf.jgss.GSSContext;
import org.ietf.jgss.GSSCredential;
import org.ietf.jgss.GSSException;
import org.ietf.jgss.MessageProp;

/**
 * This class wraps a GSSContext and provides some functionality to obtain and validate spnego tokens.
 */
public class SpnegoTokenContext {

    private static final org.slf4j.Logger LOG =
        org.slf4j.LoggerFactory.getLogger(SpnegoTokenContext.class);

    private GSSContext secContext;
    private byte[] token;
    private boolean mutualAuth;
    private SpnegoClientAction clientAction;
    private SpnegoServiceAction serviceAction;
    private GSSCredential delegationCredential;
    private Principal spnegoPrincipal;

    /**
     * Retrieve a service ticket from a KDC using the Kerberos JAAS module, and set it in this
     * BinarySecurityToken.
     * @param jaasLoginModuleName the JAAS Login Module name to use
     * @param callbackHandler a CallbackHandler instance to retrieve a password (optional)
     * @param serviceName the desired Kerberized service
     * @throws WSSecurityException
     */
    public void retrieveServiceTicket(
        String jaasLoginModuleName,
        CallbackHandler callbackHandler,
        String serviceName
    ) throws WSSecurityException {
        retrieveServiceTicket(jaasLoginModuleName, callbackHandler, serviceName, false);
    }


    /**
     * Retrieve a service ticket from a KDC using the Kerberos JAAS module, and set it in this
     * BinarySecurityToken.
     * @param jaasLoginModuleName the JAAS Login Module name to use
     * @param callbackHandler a CallbackHandler instance to retrieve a password (optional)
     * @param serviceName the desired Kerberized service
     * @param isUsernameServiceNameForm
     * @throws WSSecurityException
     */
    public void retrieveServiceTicket(
        String jaasLoginModuleName,
        CallbackHandler callbackHandler,
        String serviceName,
        boolean isUsernameServiceNameForm
    ) throws WSSecurityException {
        retrieveServiceTicket(jaasLoginModuleName, callbackHandler, serviceName,
                              isUsernameServiceNameForm, false, null);
    }

    /**
     * Retrieve a service ticket from a KDC using the Kerberos JAAS module, and set it in this
     * BinarySecurityToken.
     * @param jaasLoginModuleName the JAAS Login Module name to use
     * @param callbackHandler a CallbackHandler instance to retrieve a password (optional)
     * @param serviceName the desired Kerberized service
     * @param isUsernameServiceNameForm
     * @param requestCredDeleg Whether to request credential delegation or not
     * @param delegationCredential The delegation credential to use
     * @throws WSSecurityException
     */
    public void retrieveServiceTicket(
        String jaasLoginModuleName,
        CallbackHandler callbackHandler,
        String serviceName,
        boolean isUsernameServiceNameForm,
        boolean requestCredDeleg,
        GSSCredential delegationCredential
    ) throws WSSecurityException {

        // Get a TGT from the KDC using JAAS
        LoginContext loginContext = null;
        try {
            if (callbackHandler == null) {
                loginContext = new LoginContext(jaasLoginModuleName);
            } else {
                loginContext = new LoginContext(jaasLoginModuleName, callbackHandler);
            }
            loginContext.login();
        } catch (LoginException ex) {
            LOG.debug(ex.getMessage(), ex);
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE, ex, "kerberosLoginError",
                new Object[] {ex.getMessage()});
        }
        LOG.debug("Successfully authenticated to the TGT");

        Subject clientSubject = loginContext.getSubject();
        Set<Principal> clientPrincipals = clientSubject.getPrincipals();
        if (clientPrincipals.isEmpty()) {
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE,
                "kerberosLoginError",
                new Object[] {"No Client principals found after login"});
        }

        // Get the service ticket
        if (clientAction != null) {
            clientAction.setServiceName(serviceName);
            clientAction.setMutualAuth(mutualAuth);
            clientAction.setUserNameServiceForm(isUsernameServiceNameForm);
            token = Subject.doAs(clientSubject, clientAction);
            if (token == null) {
                throw new WSSecurityException(
                    WSSecurityException.ErrorCode.FAILURE, "kerberosServiceTicketError"
                );
            }

            secContext = clientAction.getContext();
        } else {
            KerberosClientExceptionAction action =
                new KerberosClientExceptionAction(null, serviceName,
                                                  isUsernameServiceNameForm,
                                                  requestCredDeleg,
                                                  delegationCredential,
                                                  true,
                                                  mutualAuth);
            KerberosContext krbCtx = null;
            try {
                krbCtx = (KerberosContext) Subject.doAs(clientSubject, action);

                token = krbCtx.getKerberosToken();
                if (token == null) {
                    throw new WSSecurityException(
                        WSSecurityException.ErrorCode.FAILURE, "kerberosServiceTicketError"
                    );
                }

                secContext = krbCtx.getGssContext();
            } catch (PrivilegedActionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof WSSecurityException) {
                    throw (WSSecurityException) cause;
                } else {
                    throw new WSSecurityException(
                         ErrorCode.FAILURE, new Exception(cause), "kerberosServiceTicketError"
                    );
                }
            }
        }

        LOG.debug("Successfully retrieved a service ticket");
    }

    /**
     * Validate a service ticket.
     * @param jaasLoginModuleName
     * @param callbackHandler
     * @param serviceName
     * @param ticket
     * @throws WSSecurityException
     */
    public void validateServiceTicket(
        String jaasLoginModuleName,
        CallbackHandler callbackHandler,
        String serviceName,
        byte[] ticket
    ) throws WSSecurityException {
        validateServiceTicket(jaasLoginModuleName, callbackHandler, serviceName, false, ticket);
     }

    /**
     * Validate a service ticket.
     * @param jaasLoginModuleName
     * @param callbackHandler
     * @param serviceName
     * @param ticket
     * @throws WSSecurityException
     */
    public void validateServiceTicket(
        String jaasLoginModuleName,
        CallbackHandler callbackHandler,
        String serviceName,
        boolean isUsernameServiceNameForm,
        byte[] ticket
    ) throws WSSecurityException {
        // Get a TGT from the KDC using JAAS
        LoginContext loginContext = null;
        try {
            if (callbackHandler == null) {
                loginContext = new LoginContext(jaasLoginModuleName);
            } else {
                loginContext = new LoginContext(jaasLoginModuleName, callbackHandler);
            }
            loginContext.login();
        } catch (LoginException ex) {
            LOG.debug(ex.getMessage(), ex);
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE, ex, "kerberosLoginError",
                new Object[] {ex.getMessage()});
        }
        LOG.debug("Successfully authenticated to the TGT");

        // Get the service name to use - fall back on the principal
        Subject subject = loginContext.getSubject();
        String service = serviceName;
        if (service == null) {
            Set<Principal> principals = subject.getPrincipals();
            if (principals.isEmpty()) {
                throw new WSSecurityException(
                    WSSecurityException.ErrorCode.FAILURE,
                    "kerberosLoginError",
                    new Object[] {"No Client principals found after login"});
            }
            service = principals.iterator().next().getName();
        }

        // Validate the ticket
        if (serviceAction != null) {
            serviceAction.setTicket(ticket);
            serviceAction.setServiceName(service);
            serviceAction.setUsernameServiceNameForm(isUsernameServiceNameForm);
            token = Subject.doAs(subject, serviceAction);
            secContext = serviceAction.getContext();
        } else {
            KerberosServiceExceptionAction action =
                new KerberosServiceExceptionAction(ticket, service,
                                                   isUsernameServiceNameForm, true);
            KerberosServiceContext krbCtx = null;
            try {
                krbCtx = (KerberosServiceContext) Subject.doAs(subject, action);

                token = krbCtx.getKerberosToken();
                if (token == null) {
                    throw new WSSecurityException(
                        WSSecurityException.ErrorCode.FAILURE, "kerberosServiceTicketError"
                    );
                }

                secContext = krbCtx.getGssContext();
                delegationCredential = krbCtx.getDelegationCredential();
                spnegoPrincipal = krbCtx.getPrincipal();
            } catch (PrivilegedActionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof WSSecurityException) {
                    throw (WSSecurityException) cause;
                } else {
                    throw new WSSecurityException(
                         ErrorCode.FAILURE, new Exception(cause), "kerberosServiceTicketError"
                    );
                }
            }
        }

        LOG.debug("Successfully validated a service ticket");
    }

    /**
     * Whether to enable mutual authentication or not. This only applies to retrieve service ticket.
     */
    public void setMutualAuth(boolean mutualAuthentication) {
        mutualAuth = mutualAuthentication;
    }

    /**
     * Get the SPNEGO token that was created.
     */
    public byte[] getToken() {
        return token;
    }

    /**
     * Whether a connection has been established (at the service side)
     */
    public boolean isEstablished() {
        if (secContext == null) {
            return false;
        }
        return secContext.isEstablished();
    }

    /**
     * Unwrap a key
     */
    public byte[] unwrapKey(byte[] secret) throws WSSecurityException {
        MessageProp mProp = new MessageProp(0, true);
        try {
            return secContext.unwrap(secret, 0, secret.length, mProp);
        } catch (GSSException e) {
            LOG.debug("Error in cleaning up a GSS context", e);
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE, e, "spnegoKeyError"
            );
        }
    }

    /**
     * Wrap a key
     */
    public byte[] wrapKey(byte[] secret) throws WSSecurityException {
        MessageProp mProp = new MessageProp(0, true);
        try {
            return secContext.wrap(secret, 0, secret.length, mProp);
        } catch (GSSException e) {
            LOG.debug("Error in cleaning up a GSS context", e);
            throw new WSSecurityException(
                WSSecurityException.ErrorCode.FAILURE, e, "spnegoKeyError"
            );
        }
    }

    /**
     * Set a custom SpnegoClientAction implementation to use
     */
    public void setSpnegoClientAction(SpnegoClientAction spnegoClientAction) {
        this.clientAction = spnegoClientAction;
    }

    /**
     * Set a custom SpnegoServiceAction implementation to use
     */
    public void setSpnegoServiceAction(SpnegoServiceAction spnegoServiceAction) {
        this.serviceAction = spnegoServiceAction;
    }

    public void clear() {
        token = null;
        mutualAuth = false;
        delegationCredential = null;
        spnegoPrincipal = null;
        try {
            secContext.dispose();
        } catch (GSSException e) {
            LOG.debug("Error in cleaning up a GSS context", e);
        }
    }

    public GSSCredential getDelegationCredential() {
        return delegationCredential;
    }

    public Principal getSpnegoPrincipal() {
        return spnegoPrincipal;
    }

}