/*
 * 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.atlas.security;

import com.sun.jersey.api.client.config.DefaultClientConfig;
import com.sun.jersey.client.urlconnection.HttpURLConnectionFactory;
import com.sun.jersey.client.urlconnection.URLConnectionClientHandler;
import org.apache.atlas.AtlasException;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.alias.CredentialProviderFactory;
import org.apache.hadoop.security.authentication.client.ConnectionConfigurator;
import org.apache.hadoop.security.ssl.SSLFactory;
import org.apache.hadoop.security.token.delegation.web.DelegationTokenAuthenticatedURL;
import org.apache.hadoop.security.token.delegation.web.DelegationTokenAuthenticator;
import org.apache.hadoop.security.token.delegation.web.KerberosDelegationTokenAuthenticator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.HostnameVerifier;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLSocketFactory;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.URL;
import java.net.URLConnection;
import java.security.GeneralSecurityException;
import java.security.PrivilegedExceptionAction;

import static org.apache.atlas.security.SecurityProperties.CERT_STORES_CREDENTIAL_PROVIDER_PATH;
import static org.apache.atlas.security.SecurityProperties.CLIENT_AUTH_KEY;
import static org.apache.atlas.security.SecurityProperties.KEYSTORE_FILE_KEY;
import static org.apache.atlas.security.SecurityProperties.TRUSTSTORE_FILE_KEY;

/**
 *
 */
public class SecureClientUtils {
    private static final Logger LOG = LoggerFactory.getLogger(SecureClientUtils.class);

    public static final int DEFAULT_SOCKET_TIMEOUT_IN_MSECS = 1 * 60 * 1000; // 1 minute

    private static final ConnectionConfigurator DEFAULT_TIMEOUT_CONN_CONFIGURATOR = conn -> {
        setTimeouts(conn, DEFAULT_SOCKET_TIMEOUT_IN_MSECS);

        return conn;
    };

    private SSLFactory factory;

    public static void persistSSLClientConfiguration(org.apache.commons.configuration.Configuration clientConfig, String confLocation) throws AtlasException, IOException {
        //trust settings
        Configuration configuration = new Configuration(false);
        File          sslClientFile = getSSLClientFile(confLocation);

        if (!sslClientFile.exists()) {
            configuration.set("ssl.client.truststore.type", "jks");
            configuration.set("ssl.client.truststore.location", clientConfig.getString(TRUSTSTORE_FILE_KEY));

            if (clientConfig.getBoolean(CLIENT_AUTH_KEY, false)) {
                // need to get client key properties
                configuration.set("ssl.client.keystore.location", clientConfig.getString(KEYSTORE_FILE_KEY));
                configuration.set("ssl.client.keystore.type", "jks");
            }

            // add the configured credential provider
            configuration.set(CredentialProviderFactory.CREDENTIAL_PROVIDER_PATH, clientConfig.getString(CERT_STORES_CREDENTIAL_PROVIDER_PATH));

            String hostnameVerifier = clientConfig.getString(SSLFactory.SSL_HOSTNAME_VERIFIER_KEY);

            if (hostnameVerifier != null) {
                configuration.set(SSLFactory.SSL_HOSTNAME_VERIFIER_KEY, hostnameVerifier);
            }

            configuration.writeXml(new FileWriter(sslClientFile));
        }
    }

    public URLConnectionClientHandler getClientConnectionHandler(DefaultClientConfig config, org.apache.commons.configuration.Configuration clientConfig, String doAsUser, final UserGroupInformation ugi) {
        config.getProperties().put(URLConnectionClientHandler.PROPERTY_HTTP_URL_CONNECTION_SET_METHOD_WORKAROUND, true);

        Configuration conf = new Configuration();

        conf.addResource(conf.get(SSLFactory.SSL_CLIENT_CONF_KEY, SecurityProperties.SSL_CLIENT_PROPERTIES));

        UserGroupInformation.setConfiguration(conf);

        final ConnectionConfigurator connConfigurator = newConnConfigurator(conf);
        DelegationTokenAuthenticator authenticator    = new KerberosDelegationTokenAuthenticator();

        authenticator.setConnectionConfigurator(connConfigurator);

        final DelegationTokenAuthenticator          finalAuthenticator       = authenticator;
        final DelegationTokenAuthenticatedURL.Token token                    = new DelegationTokenAuthenticatedURL.Token();
        HttpURLConnectionFactory                    httpURLConnectionFactory = null;

        try {
            UserGroupInformation       ugiToUse  = ugi != null ? ugi : UserGroupInformation.getCurrentUser();
            final UserGroupInformation actualUgi = (ugiToUse.getAuthenticationMethod() == UserGroupInformation.AuthenticationMethod.PROXY)                     ? ugiToUse.getRealUser() : ugiToUse;

            LOG.info("Real User: {}, is from ticket cache? {}", actualUgi, UserGroupInformation.isLoginTicketBased());

            if (StringUtils.isEmpty(doAsUser) || StringUtils.equals(doAsUser, actualUgi.getShortUserName())) {
                doAsUser = null;
            }

            LOG.info("doAsUser: {}", doAsUser);

            final String finalDoAsUser = doAsUser;

            httpURLConnectionFactory = url -> {
                try {
                    return actualUgi.doAs((PrivilegedExceptionAction<HttpURLConnection>) () -> {
                        try {
                            return new DelegationTokenAuthenticatedURL(finalAuthenticator, connConfigurator).openConnection(url, token, finalDoAsUser);
                        } catch (Exception e) {
                            throw new IOException(e);
                        }
                    });
                } catch (Exception e) {
                    if (e instanceof IOException) {
                        throw (IOException) e;
                    } else {
                        throw new IOException(e);
                    }
                }
            };
        } catch (IOException e) {
            LOG.warn("Error obtaining user", e);
        }

        return new URLConnectionClientHandler(httpURLConnectionFactory);
    }

    public SSLFactory getSSLFactory(Configuration conf) throws IOException, GeneralSecurityException {
        if (factory == null) {
            factory = new SSLFactory(SSLFactory.Mode.CLIENT, conf);

            factory.init();
        }

        return factory;
    }

    public void destroyFactory() {
        if (factory != null) {
            factory.destroy();

            factory = null;
        }
    }

    public URLConnectionClientHandler getUrlConnectionClientHandler() {
        return new URLConnectionClientHandler(url -> {
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();

            if (connection instanceof HttpsURLConnection) {
                LOG.debug("Attempting to configure HTTPS connection using client configuration");

                try {
                    Configuration conf = new Configuration();

                    conf.addResource(conf.get(SSLFactory.SSL_CLIENT_CONF_KEY, SecurityProperties.SSL_CLIENT_PROPERTIES));

                    UserGroupInformation.setConfiguration(conf);

                    HttpsURLConnection c = (HttpsURLConnection) connection;

                    final SSLFactory       factory = getSSLFactory(conf);
                    final SSLSocketFactory sf      = factory.createSSLSocketFactory();
                    final HostnameVerifier hv      = factory.getHostnameVerifier();

                    c.setSSLSocketFactory(sf);
                    c.setHostnameVerifier(hv);
                } catch (Exception e) {
                    LOG.info("Unable to configure HTTPS connection from configuration. Leveraging JDK properties.");
                }
            }

            return connection;
        });
    }

    private ConnectionConfigurator newConnConfigurator(Configuration conf) {
        try {
            return newSslConnConfigurator(DEFAULT_SOCKET_TIMEOUT_IN_MSECS, conf);
        } catch (Exception e) {
            LOG.debug("Cannot load customized ssl related configuration. Fallback to system-generic settings.", e);

            return DEFAULT_TIMEOUT_CONN_CONFIGURATOR;
        }
    }

    private ConnectionConfigurator newSslConnConfigurator(final int timeout, Configuration conf) throws IOException, GeneralSecurityException {
        final SSLFactory       factory = getSSLFactory(conf);
        final SSLSocketFactory sf      = factory.createSSLSocketFactory();
        final HostnameVerifier hv      = factory.getHostnameVerifier();

        return conn -> {
            if (conn instanceof HttpsURLConnection) {
                HttpsURLConnection c = (HttpsURLConnection) conn;

                c.setSSLSocketFactory(sf);
                c.setHostnameVerifier(hv);
            }

            setTimeouts(conn, timeout);

            return conn;
        };
    }

    private static void setTimeouts(URLConnection connection, int socketTimeout) {
        connection.setConnectTimeout(socketTimeout);
        connection.setReadTimeout(socketTimeout);
    }

    private static File getSSLClientFile(String confLocation) throws AtlasException {
        File sslDir;
        try {
            if (confLocation == null) {
                String persistDir = null;
                URL    resource   = SecureClientUtils.class.getResource("/");

                if (resource != null) {
                    persistDir = resource.toURI().getPath();
                }

                sslDir = new File(persistDir);
            } else {
                sslDir = new File(confLocation);
            }

            LOG.info("ssl-client.xml will be created in {}", sslDir);
        } catch (Exception e) {
            throw new AtlasException("Failed to find client configuration directory", e);
        }

        return new File(sslDir, SecurityProperties.SSL_CLIENT_PROPERTIES);
    }
}
