/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2016 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * Licensed 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.wildfly.security.mechanism.gssapi;

import static javax.security.auth.login.AppConfigurationEntry.LoginModuleControlFlag.REQUIRED;
import static org.wildfly.common.Assert.checkNotNullParam;
import static org.wildfly.security.mechanism.gssapi.ElytronMessages.log;

import java.io.File;
import java.io.IOException;
import java.security.AccessController;
import java.security.GeneralSecurityException;
import java.security.PrivilegedAction;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.UnaryOperator;

import javax.security.auth.RefreshFailedException;
import javax.security.auth.Subject;
import javax.security.auth.kerberos.KerberosPrincipal;
import javax.security.auth.kerberos.KerberosTicket;
import javax.security.auth.kerberos.KeyTab;
import javax.security.auth.login.AppConfigurationEntry;
import javax.security.auth.login.Configuration;
import javax.security.auth.login.LoginContext;
import javax.security.auth.login.LoginException;

import org.ietf.jgss.GSSCredential;
import org.ietf.jgss.GSSException;
import org.ietf.jgss.GSSManager;
import org.ietf.jgss.GSSName;
import org.ietf.jgss.Oid;
import org.wildfly.common.function.ExceptionSupplier;
import org.wildfly.security.SecurityFactory;
import org.wildfly.security.auth.callback.FastUnsupportedCallbackException;
import org.wildfly.security.credential.GSSKerberosCredential;
import org.wildfly.security.manager.action.SetContextClassLoaderAction;

/**
 * A {@link SecurityFactory} implementation for obtaining a {@link GSSCredential}.
 *
 * @author <a href="mailto:darran.lofthouse@jboss.com">Darran Lofthouse</a>
 */
public final class GSSCredentialSecurityFactory implements SecurityFactory<GSSKerberosCredential> {

    private static final String KRB5LoginModule = "com.sun.security.auth.module.Krb5LoginModule";
    private static final long ONE_SECOND = 1000;

    public static final Oid KERBEROS_V5;
    public static final Oid SPNEGO;

    static {
        try {
            KERBEROS_V5 = new Oid("1.2.840.113554.1.2.2");
            SPNEGO = new Oid("1.3.6.1.5.5.2");
        } catch (GSSException e) {
            throw new RuntimeException("Unable to initialise Oid", e);
        }
    }

    private final int minimumRemainingLifetime;
    private final ExceptionSupplier<GSSKerberosCredential, GeneralSecurityException> rawSupplier;

    private final AtomicReference<GSSKerberosCredential> cachedCredentialReference = new AtomicReference<>();
    private final UnaryOperator<GSSKerberosCredential> credentialOperator;


    /**
     * Constructs a new {@code GSSCredentialSecurityFactory} instance.
     *
     * @param minimumRemainingLifetime the minimum remaining lifetime for a {@link GSSCredential} in seconds.
     * @param rawSupplier the supplier of raw credentials.
     */
    GSSCredentialSecurityFactory(final int minimumRemainingLifetime, final ExceptionSupplier<GSSKerberosCredential, GeneralSecurityException> rawSupplier) {
        this.minimumRemainingLifetime = minimumRemainingLifetime;
        this.rawSupplier = rawSupplier;
        credentialOperator = this::update;
    }

    /**
     * Updates the {@link GSSKerberosCredential}. If the original is not valid, it gets a new {@code GSSKerberosCredential}
     * from the {@code rawSupplier}, otherwise returns the original.
     *
     * @param original the original {@code GSSKerberosCredential} to be updated.
     * @return the original if still valid, new {@code GSSKerberosCredential} otherwise.
     */
    private GSSKerberosCredential update(GSSKerberosCredential original) {
        GSSKerberosCredential result = null;
        try {
            if (original != null) {
                if (testIsValid(original.getGssCredential()) && testIsValid(original.getKerberosTicket())) {
                    result = original;
                }
            }

            if (result == null) {
                log.trace("No valid cached credential, obtaining new one...");
                result = rawSupplier.get();
                log.tracef("Obtained GSSCredentialCredential [%s]", result);
            } else {
                log.tracef("Used cached GSSCredential [%s]", result);
            }
        } catch (GeneralSecurityException e) {
            throw new IllegalStateException(e);
        }

        return result;
    }

    /**
     * Checks if the GSSCredential is still valid.
     *
     * @param gssCredential the GSSCredential to check.
     * @return {@code true} if the GSSCredential is valid, {@code false} otherwise.
     * @throws GeneralSecurityException if an error occurs during the validation.
     */
    private boolean testIsValid(GSSCredential gssCredential) throws GeneralSecurityException {
        checkNotNullParam("gssCredential", gssCredential);
        boolean stillValid;
        try {
            int remainingLifetime = gssCredential.getRemainingLifetime();
            log.tracef("Remaining GSSCredential Lifetime = %d", remainingLifetime);
            stillValid = remainingLifetime >= minimumRemainingLifetime;
        } catch (GSSException e) {
            throw new GeneralSecurityException(e);
        }

        log.tracef("testIsValid(GSSCredential)=%b", stillValid);
        return stillValid;
    }

    /**
     * Checks if the Kerberos ticket is still valid. If not, attempts to refresh it.
     *
     * @param ticket the Kerberos ticket to be checked.
     * @return {@code true} if the ticket is valid, {@code false} otherwise.
     */
    private boolean testIsValid(KerberosTicket ticket) {
        if (ticket == null) {
            log.trace("No cached KerberosTicket");
            return true; // If there is no ticket it is not "invalid".
        }

        Date endTime = ticket.getEndTime();
        log.tracef("KerberosTicket.getEndTime()=%s", endTime);
        boolean stillValid = endTime != null && System.currentTimeMillis() < endTime.getTime() - (minimumRemainingLifetime * ONE_SECOND);

        if (!stillValid) {
            log.trace("Attempting to refresh existing KerberosTicket.");
            try {
                ticket.refresh();
                log.tracef("KerberosTicket refreshed until %s", ticket.getEndTime());
                stillValid = true;
            } catch (RefreshFailedException e) {
                log.tracef("Unable to refresh KerberosTicket.", e);
            }
        }

        log.tracef("testIsValid(KerberosTicket)=%b", stillValid);
        return stillValid;
    }

    @Override
    public GSSKerberosCredential create() throws GeneralSecurityException {
        try {
            return cachedCredentialReference.updateAndGet(credentialOperator);
        } catch (RuntimeException e) {
            if (e.getCause() instanceof GSSException) {
                throw new GeneralSecurityException(e.getCause());
            } else if (e.getCause() instanceof GeneralSecurityException) {
                throw (GeneralSecurityException) e.getCause();
            }

            throw e;
        }
    }


    /**
     * Obtain a new {@link Builder} capable of building a {@link GSSCredentialSecurityFactory}.
     *
     * @return a new {@link Builder} capable of building a {@link GSSCredentialSecurityFactory}.
     */
    public static Builder builder() {
        return new Builder();
    }

    /**
     * A builder for GSS credential security factories.
     */
    public static class Builder {

        private boolean built = false;
        private List<Oid> mechanismOids = new ArrayList<>();
        private String principal;
        private File keyTab;
        private boolean isServer;
        private boolean obtainKerberosTicket;
        private int minimumRemainingLifetime;
        private int requestLifetime;
        private boolean debug;
        private boolean wrapGssCredential;
        private boolean checkKeyTab;
        private volatile long lastFailTime = 0;
        private long failCache = 0;
        private Map<String, Object> options;

        Builder() {
        }

        /**
         * Set the keytab file to obtain the identity.
         *
         * @param keyTab the keytab file to obtain the identity.
         * @return {@code this} to allow chaining.
         */
        public Builder setKeyTab(final File keyTab) {
            assertNotBuilt();
            this.keyTab = keyTab;

            return this;
        }

        /**
         * Set if the credential returned from the factory is representing the server side of the connection.
         *
         * @param isServer is the credential returned from the factory is representing the server side of the connection.
         * @return {@code this} to allow chaining.
         */
        public Builder setIsServer(final boolean isServer) {
            assertNotBuilt();
            this.isServer = isServer;

            return this;
        }

        /**
         * Set if the KerberosTicket should also be obtained and associated with the Credential.
         *
         * @param obtainKerberosTicket if the KerberosTicket should also be obtained and associated with the Credential.
         * @return {@code this} to allow chaining.
         */
        public Builder setObtainKerberosTicket(final boolean obtainKerberosTicket) {
            assertNotBuilt();
            this.obtainKerberosTicket = obtainKerberosTicket;

            return this;
        }

        /**
         * Once the factory has been called once it will cache the resulting {@link GSSCredential}, this setting
         * defines how much life it must have left in seconds for it to be re-used.
         *
         * @param minimumRemainingLifetime the time in seconds of life a {@link GSSCredential} must have to be re-used.
         * @return {@code this} to allow chaining.
         */
        public Builder setMinimumRemainingLifetime(final int minimumRemainingLifetime) {
            assertNotBuilt();
            this.minimumRemainingLifetime = minimumRemainingLifetime;

            return this;
        }

        /**
         * Set the lifetime to request newly created credentials are valid for.
         *
         * @param requestLifetime the lifetime to request newly created credentials are valid for.
         * @return {@code this} to allow chaining.
         */
        public Builder setRequestLifetime(final int requestLifetime) {
            assertNotBuilt();
            this.requestLifetime = requestLifetime < 0 ? GSSCredential.INDEFINITE_LIFETIME : requestLifetime;

            return this;
        }

        /**
         * Add an {@link Oid} for a mechanism the {@link GSSCredential} should be usable with.
         *
         * @param oid the {@link Oid} for the mechanism the {@link GSSCredential} should be usable with.
         * @return {@code this} to allow chaining.
         */
        public Builder addMechanismOid(final Oid oid) {
            assertNotBuilt();
            mechanismOids.add(checkNotNullParam("oid", oid));

            return this;
        }

        /**
         * Set the principal name for the initial authentication from the KeyTab.
         *
         * @param principal the principal name for the initial authentication from the KeyTab.
         * @return {@code this} to allow chaining.
         */
        public Builder setPrincipal(final String principal) {
            assertNotBuilt();
            this.principal = principal;

            return this;
        }

        /**
         * Set if debug logging should be enabled for the JAAS authentication portion of obtaining the {@link GSSCredential}.
         *
         * @param debug if debug logging should be enabled for the JAAS authentication portion of obtaining the {@link GSSCredential}
         * @return {@code this} to allow chaining.
         */
        public Builder setDebug(final boolean debug) {
            assertNotBuilt();
            this.debug = debug;

            return this;
        }

        /**
         * Set if the constructed {@link GSSCredential} should be wrapped to prevent improper credential disposal or not.
         *
         * @param value {@code true} if the constructed {@link GSSCredential} should be wrapped; {@code false} otherwise.
         * @return {@code this} to allow chaining.
         */
        public Builder setWrapGssCredential(final boolean value) {
            assertNotBuilt();
            this.wrapGssCredential = value;

            return this;
        }

        /**
         * Set if keytab file existence and principal presence in it should be checked on factory build.
         *
         * @param value {@code true} if keytab file should be checked; {@code false} otherwise.
         * @return {@code this} to allow chaining.
         */
        public Builder setCheckKeyTab(final boolean value) {
            assertNotBuilt();
            this.checkKeyTab = value;

            return this;
        }

        /**
         * Set other configuration options for {@code Krb5LoginModule}.
         *
         * @param options the configuration options which will be appended to options passed into {@code Krb5LoginModule}
         * @return {@code this} to allow chaining.
         */
        public Builder setOptions(final Map<String, Object> options) {
            assertNotBuilt();
            this.options = options;

            return this;
        }

        /**
         * Set amount of seconds before new try to obtain {@link GSSCredential} should be done if it has failed last time.
         * Allows to prevent long waiting to unavailable KDC on every authentication.
         *
         * @param seconds amount of seconds to cache fail state of the credential factory; 0 if the cache should not be used.
         * @return {@code this} to allow chaining.
         */
        public Builder setFailCache(final long seconds) {
            assertNotBuilt();
            this.failCache = seconds;

            return this;
        }

        /**
         * Construct a new {@link GSSKerberosCredential} security factory instance.
         *
         * @return the built factory instance
         * @throws IOException when unable to use given KeyTab
         */
        public SecurityFactory<GSSKerberosCredential> build() throws IOException {
            assertNotBuilt();
            if (checkKeyTab) {
                checkKeyTab();
            }

            final Configuration configuration = createConfiguration();

            built = true;
            return new GSSCredentialSecurityFactory(minimumRemainingLifetime > 0 ? minimumRemainingLifetime : 0, () -> createGSSCredential(configuration));
        }

        /**
         * Creates an instance of the {@link GSSKerberosCredential} class, which represents a Kerberos credential
         * that can be used for authentication using the GSS-API.
         *
         * @param configuration the configuration used for creating the {@link LoginContext}.
         * @return the {@code GSSKerberosCredential} - the GSSCredential object and Kerberos Ticket (if {@code obtainKerberosTicket} is {@code true}.
         * @throws GeneralSecurityException if an error occurs during the creation of {@code GSSKerberosCredential}.
         */
        private GSSKerberosCredential createGSSCredential(Configuration configuration) throws GeneralSecurityException {
            if (failCache != 0 && System.currentTimeMillis() - lastFailTime < failCache * 1000) {
                throw log.initialLoginSkipped(failCache);
            }

            final Subject subject = new Subject();

            try {
                final ClassLoader oldCl = doPrivileged(new SetContextClassLoaderAction(Builder.class.getClassLoader()));
                final LoginContext lc;
                try {
                    lc = new LoginContext("KDC", subject, (c) -> {
                        throw new FastUnsupportedCallbackException(c[0]);
                    }, configuration);
                } finally {
                    doPrivileged(new SetContextClassLoaderAction(oldCl));
                }
                log.tracef("Logging in using LoginContext and subject [%s]", subject);
                lc.login();
                log.tracef("Logging in using LoginContext and subject [%s] succeed", subject);

                final KerberosTicket kerberosTicket;
                if (obtainKerberosTicket) {
                    Set<KerberosTicket> kerberosTickets = doPrivileged((PrivilegedAction<Set<KerberosTicket>>) () -> subject.getPrivateCredentials(KerberosTicket.class));
                    if (kerberosTickets.size() > 1) {
                        throw log.tooManyKerberosTicketsFound();
                    }
                    kerberosTicket = kerberosTickets.size() == 1 ? kerberosTickets.iterator().next() : null;
                } else {
                    kerberosTicket = null;
                }

                final GSSManager manager = GSSManager.getInstance();
                return Subject.doAs(subject, (PrivilegedExceptionAction<GSSKerberosCredential>) () -> {
                    Set<KerberosPrincipal> principals = subject.getPrincipals(KerberosPrincipal.class);
                    if (principals.size() < 1) {
                        throw log.noKerberosPrincipalsFound();
                    } else if (principals.size() > 1) {
                        throw log.tooManyKerberosPrincipalsFound();
                    }
                    KerberosPrincipal principal = principals.iterator().next();
                    log.tracef("Creating GSSName for Principal '%s'", principal);
                    GSSName name = manager.createName(principal.getName(), GSSName.NT_USER_NAME, KERBEROS_V5);

                    if (wrapGssCredential) {
                        return new GSSKerberosCredential(wrapCredential(manager.createCredential(name, requestLifetime, mechanismOids.toArray(new Oid[mechanismOids.size()]),
                                isServer ? GSSCredential.ACCEPT_ONLY : GSSCredential.INITIATE_ONLY)), kerberosTicket);
                    }
                    return new GSSKerberosCredential(manager.createCredential(name, requestLifetime, mechanismOids.toArray(new Oid[mechanismOids.size()]),
                            isServer ? GSSCredential.ACCEPT_ONLY : GSSCredential.INITIATE_ONLY), kerberosTicket);
                });

            } catch (LoginException e) {
                if (failCache != 0) {
                    lastFailTime = System.currentTimeMillis();
                }
                throw log.unableToPerformInitialLogin(e);
            } catch (PrivilegedActionException e) {
                if (e.getCause() instanceof GeneralSecurityException) {
                    throw (GeneralSecurityException) e.getCause();
                }
                throw new GeneralSecurityException(e.getCause());
            }
        }

        /**
         * Performs a privileged action. If a security manager is set, the action will be executed via
         * {@link AccessController#doPrivileged(PrivilegedAction)}. If no security manager is set,
         * the action will be executed directly.
         *
         * @param action the action do be executed.
         * @param <T> the type of the action.
         * @return the result of the executed action.
         */
        private static <T> T doPrivileged(final PrivilegedAction<T> action) {
            return System.getSecurityManager() != null ? AccessController.doPrivileged(action) : action.run();
        }

        /**
         * Checks if the keytab exists and if it contains any keys for the specified principal.
         *
         * @throws IOException if the keytab does not exist or if it does not contain any keys for the specified principal.
         */
        private void checkKeyTab() throws IOException {
            KeyTab kt = KeyTab.getInstance(keyTab);
            if (!kt.exists()) {
                throw log.keyTabDoesNotExists(keyTab.getAbsolutePath());
            }
            if (kt.getKeys(new KerberosPrincipal(principal)).length == 0) {
                throw log.noKeysForPrincipalInKeyTab(principal, keyTab.getAbsolutePath());
            }
        }

        /**
         * Creates a {@link Configuration} that is used to initiate a {@link LoginContext}.
         *
         * @return a {@code Configuration} for initiating a {@code LoginContext}.
         * @throws IOException if the keyTab does not exist or there are no keys for the principal in the keyTab.
         */
        private Configuration createConfiguration() throws IOException {
            Map<String, Object> options = new HashMap<>();
            if (debug) {
                options.put("debug", "true");
            }
            options.put("principal", principal);
            options.put("storeKey", "true");
            options.put("useKeyTab", "true");
            if (keyTab != null) options.put("keyTab", keyTab.getAbsolutePath());
            options.put("isInitiator", (isServer && !obtainKerberosTicket) ? "false" : "true");

            if (this.options != null) {
                options.putAll(this.options);
            }

            log.tracef("Created LoginContext configuration: %s", options.toString());

            final AppConfigurationEntry[] aceArray = new AppConfigurationEntry[] {
                    new AppConfigurationEntry(KRB5LoginModule, REQUIRED, options)
            };

            return new Configuration() {

                @Override
                public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
                    assert "KDC".equals(name);
                    return aceArray;
                }

            };
        }

        /**
         * Asserts that the builder has not yet been built.
         */
        private void assertNotBuilt() {
            if (built) {
                throw log.builderAlreadyBuilt();
            }
        }

    }

    /**
     * Wraps the given {@link GSSCredential} and prevents it from being disposed.
     *
     * @param credential the {@code GSSCredential} to be wrapped.
     * @return the wrapped {@code GSSCredential}.
     */
    private static GSSCredential wrapCredential(final GSSCredential credential) {
        return new GSSCredential() {

            @Override
            public int getUsage(Oid mech) throws GSSException {
                return credential.getUsage(mech);
            }

            @Override
            public int getUsage() throws GSSException {
                return credential.getUsage();
            }

            @Override
            public int getRemainingLifetime() throws GSSException {
                return credential.getRemainingLifetime();
            }

            @Override
            public int getRemainingInitLifetime(Oid mech) throws GSSException {
                return credential.getRemainingInitLifetime(mech);
            }

            @Override
            public int getRemainingAcceptLifetime(Oid mech) throws GSSException {
                return credential.getRemainingAcceptLifetime(mech);
            }

            @Override
            public GSSName getName(Oid mech) throws GSSException {
                return credential.getName(mech);
            }

            @Override
            public GSSName getName() throws GSSException {
                return credential.getName();
            }

            @Override
            public Oid[] getMechs() throws GSSException {
                return credential.getMechs();
            }

            @Override
            public void dispose() throws GSSException {
                // Prevent disposal of our credential.
            }

            @Override
            public void add(GSSName name, int initLifetime, int acceptLifetime, Oid mech, int usage) throws GSSException {
                credential.add(name, initLifetime, acceptLifetime, mech, usage);
            }

        };
    }
}
