/*
 *
 * 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.bookkeeper.sasl;

import java.security.Principal;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
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.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.sasl.AuthorizeCallback;
import javax.security.sasl.RealmCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;
import org.apache.zookeeper.server.auth.KerberosName;
import org.slf4j.LoggerFactory;

/**
 * A SASL Client State data object.
 */
public class SaslClientState {

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

    private final SaslClient saslClient;
    private final Subject clientSubject;
    private String username;
    private String password;

    public SaslClientState(String serverHostname, Subject subject) throws SaslException {
        String saslServiceName = System.getProperty(SaslConstants.SASL_SERVICE_NAME,
                                                    SaslConstants.SASL_SERVICE_NAME_DEFAULT);
        String serverPrincipal = saslServiceName + "/" + serverHostname;
        this.clientSubject = subject;
        if (clientSubject == null) {
            throw new SaslException("Cannot create JAAS Sujbect for SASL");
        }
        if (clientSubject.getPrincipals().isEmpty()) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Using JAAS/SASL/DIGEST-MD5 auth to connect to {}", serverPrincipal);
            }
            String[] mechs = {"DIGEST-MD5"};
            username = (String) (clientSubject.getPublicCredentials().toArray()[0]);
            password = (String) (clientSubject.getPrivateCredentials().toArray()[0]);
            saslClient = Sasl.createSaslClient(mechs, username, SaslConstants.SASL_BOOKKEEPER_PROTOCOL,
                SaslConstants.SASL_MD5_DUMMY_HOSTNAME, null, new ClientCallbackHandler(password));
        } else { // GSSAPI/Kerberos
            final Object[] principals = clientSubject.getPrincipals().toArray();
            final Principal clientPrincipal = (Principal) principals[0];
            final KerberosName clientKerberosName = new KerberosName(clientPrincipal.getName());
            KerberosName serviceKerberosName = new KerberosName(serverPrincipal + "@" + clientKerberosName.getRealm());
            final String serviceName = serviceKerberosName.getServiceName();
            final String serviceHostname = serviceKerberosName.getHostName();
            final String clientPrincipalName = clientKerberosName.toString();
            if (LOG.isDebugEnabled()) {
                LOG.debug("Using JAAS/SASL/GSSAPI auth to connect to server Principal {}", serverPrincipal);
            }
            try {
                saslClient = Subject.doAs(clientSubject, new PrivilegedExceptionAction<SaslClient>() {
                    @Override
                    public SaslClient run() throws SaslException {
                        String[] mechs = {"GSSAPI"};
                        return Sasl.createSaslClient(mechs, clientPrincipalName, serviceName, serviceHostname, null,
                            new ClientCallbackHandler(null));
                    }
                });
            } catch (PrivilegedActionException err) {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("GSSAPI client error", err.getCause());
                }
                throw new SaslException("error while booting GSSAPI client", err.getCause());
            }
        }
        if (saslClient == null) {
            throw new SaslException("Cannot create JVM SASL Client");
        }

    }

    public byte[] evaluateChallenge(final byte[] saslToken) throws SaslException {
        if (saslToken == null) {
            throw new SaslException("saslToken is null");
        }
        if (clientSubject != null) {
            try {
                final byte[] retval = Subject.doAs(clientSubject, new PrivilegedExceptionAction<byte[]>() {
                        @Override
                        public byte[] run() throws SaslException {
                            return saslClient.evaluateChallenge(saslToken);
                        }
                    });
                return retval;
            } catch (PrivilegedActionException e) {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("SASL error", e.getCause());
                }
                throw new SaslException("SASL/JAAS error", e.getCause());
            }
        } else {
            return saslClient.evaluateChallenge(saslToken);
        }
    }

    public boolean hasInitialResponse() {
        return saslClient.hasInitialResponse();
    }

    static class ClientCallbackHandler implements CallbackHandler {

        private String password = null;

        public ClientCallbackHandler(String password) {
            this.password = password;
        }

        @Override
        public void handle(Callback[] callbacks) throws
            UnsupportedCallbackException {
            for (Callback callback : callbacks) {
                if (callback instanceof NameCallback) {
                    NameCallback nc = (NameCallback) callback;
                    nc.setName(nc.getDefaultName());
                } else {
                    if (callback instanceof PasswordCallback) {
                        PasswordCallback pc = (PasswordCallback) callback;
                        if (password != null) {
                            pc.setPassword(this.password.toCharArray());
                        }
                    } else {
                        if (callback instanceof RealmCallback) {
                            RealmCallback rc = (RealmCallback) callback;
                            rc.setText(rc.getDefaultText());
                        } else {
                            if (callback instanceof AuthorizeCallback) {
                                AuthorizeCallback ac = (AuthorizeCallback) callback;
                                String authid = ac.getAuthenticationID();
                                String authzid = ac.getAuthorizationID();
                                if (authid.equals(authzid)) {
                                    ac.setAuthorized(true);
                                } else {
                                    ac.setAuthorized(false);
                                }
                                if (ac.isAuthorized()) {
                                    ac.setAuthorizedID(authzid);
                                }
                            } else {
                                throw new UnsupportedCallbackException(callback, "Unrecognized SASL ClientCallback");
                            }
                        }
                    }
                }
            }
        }
    }

    public boolean isComplete() {
        return saslClient.isComplete();
    }

    public byte[] saslResponse(byte[] saslTokenMessage) {
        try {
            byte[] retval = saslClient.evaluateChallenge(saslTokenMessage);
            return retval;
        } catch (SaslException e) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("saslResponse: Failed to respond to SASL server's token:", e);
            }
            return null;
        }
    }

}
