/**
 *  Copyright 2005-2016 Red Hat, Inc.
 *
 *  Red Hat 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 io.fabric8.jaas;

import org.apache.curator.framework.CuratorFramework;
import org.apache.karaf.jaas.boot.principal.RolePolicy;
import org.apache.karaf.jaas.boot.principal.RolePrincipal;
import org.apache.karaf.jaas.boot.principal.UserPrincipal;
import org.apache.karaf.jaas.modules.BackingEngine;
import org.apache.karaf.jaas.modules.Encryption;
import org.apache.karaf.jaas.modules.encryption.EncryptionSupport;
import io.fabric8.zookeeper.curator.CuratorFrameworkLocator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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.auth.login.FailedLoginException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;

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

import static io.fabric8.zookeeper.utils.ZooKeeperUtils.getContainerTokens;
import static io.fabric8.zookeeper.utils.ZooKeeperUtils.getProperties;
import static io.fabric8.zookeeper.utils.ZooKeeperUtils.isContainerLogin;

public class ZookeeperLoginModule implements LoginModule {


    private static final Logger LOG = LoggerFactory.getLogger(ZookeeperLoginModule.class);

    private Set<Principal> principals = new HashSet<Principal>();
    private CallbackHandler callbackHandler;
    private String roleDiscriminator;
    private String rolePolicy;
    private Subject subject;
    private boolean debug = false;
    private static Properties users = new Properties();
    private static long usersTs = -1L;
    private static Properties containers = new Properties();
    private static long containersTs = -1L;
    private EncryptionSupport encryptionSupport;
    private String path;

    @Override
    @SuppressWarnings({ "rawtypes", "unchecked" })
    public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
        this.callbackHandler = callbackHandler;
        this.subject = subject;
        this.roleDiscriminator = (String) options.get("role.discriminator");
        this.rolePolicy = (String) options.get("role.policy");
        this.encryptionSupport = new BasicEncryptionSupport(options);
        this.debug = Boolean.parseBoolean((String) options.get("debug"));
        this.path = (String)options.get("path");
        if (path == null) {
            path = ZookeeperBackingEngine.USERS_NODE;
        }
        try {
            CuratorFramework curator = CuratorFrameworkLocator.getCuratorFramework();
            if (curator != null) {
                try {
                    users = getCachedUsers(curator, path, false);
                    containers = getCachedContainerTokens(curator, false);
                } catch (IllegalStateException e) {
                    if ("Client is not started".equals(e.getMessage())) {
                        LOG.warn("Zookeeper connection not available. ZK authentication module is not enabled.");
                    }
                }
            }
            if (debug) {
                LOG.debug("Initialize [" + this + "] - curator=" + curator + ",users=" + users);
            }
        } catch (Exception e) {
            LOG.warn("Failed fetching authentication data.", e);
        }
    }

    private Properties getCachedUsers(CuratorFramework curator, String path, boolean force) throws Exception {
        if (!force && usersTs + 60000L > System.currentTimeMillis()) {
            return users;
        }
        synchronized (this) {
            users = getProperties(curator, path);
            usersTs = System.currentTimeMillis();
            return users;
        }
    }

    private Properties getCachedContainerTokens(CuratorFramework curator, boolean force) throws Exception {
        if (!force && containersTs + 60000L > System.currentTimeMillis()) {
            return containers;
        }
        synchronized (this) {
            containers = getContainerTokens(curator);
            containersTs = System.currentTimeMillis();
            return containers;
        }
    }

    @Override
    public boolean login() throws LoginException {
        boolean result;
        String user = null;
        try {
            Callback[] callbacks = new Callback[2];
            callbacks[0] = new NameCallback("Username: ");
            callbacks[1] = new PasswordCallback("Password: ", false);
            try {
                callbackHandler.handle(callbacks);
            } catch (IOException ioe) {
                throw new LoginException(ioe.getMessage());
            } catch (UnsupportedCallbackException uce) {
                throw new LoginException(uce.getMessage() + " not available to obtain information from user");
            }

            user = ((NameCallback) callbacks[0]).getName();
            if (user == null)
                throw new FailedLoginException("user name is null");

            if( user.startsWith(BackingEngine.GROUP_PREFIX) ) {
                throw new IllegalArgumentException("Prefix not permitted in user names: " + BackingEngine.GROUP_PREFIX);
            }

            char[] tmpPassword = ((PasswordCallback) callbacks[1]).getPassword();
            if (tmpPassword == null) {
                tmpPassword = new char[0];
            }

            if (debug)
                LOG.debug("Login [" + this + "] - user=" + user + ",users=" + users);

            if (isContainerLogin(user)) {
                String token = containers.getProperty(user);
                if (token == null) {
                    // force reload cache of container tokens
                    CuratorFramework curator = CuratorFrameworkLocator.getCuratorFramework();
                    if (curator != null) {
                        try {
                            getCachedContainerTokens(curator, true);
                            token = containers.getProperty(user);
                        } catch (Exception e) {
                            LOG.warn(e.getMessage());
                        }
                    }
                    // didn't help
                    if (token == null) {
                        throw new FailedLoginException("Container doesn't exist");
                    }
                }

                // the password is in the first position
                if (!new String(tmpPassword).equals(token)) {
                    // force reload cache of container tokens
                    CuratorFramework curator = CuratorFrameworkLocator.getCuratorFramework();
                    if (curator != null) {
                        try {
                            getCachedContainerTokens(curator, true);
                            token = containers.getProperty(user);
                        } catch (Exception e) {
                            LOG.warn(e.getMessage());
                        }
                    }
                    // didn't help
                    if (!new String(tmpPassword).equals(token)) {
                        throw new FailedLoginException("Tokens do not match");
                    }
                }

                principals = new HashSet<Principal>();
                principals.add(new UserPrincipal(user));
                principals.add(new RolePrincipal("container"));
                principals.add(new RolePrincipal("admin"));
                subject.getPrivateCredentials().add(new String(tmpPassword));
                result = true;
            } else {
                String userInfos = users.getProperty(user);
                if (userInfos == null) {
                    // force reload cache of user tokens
                    CuratorFramework curator = CuratorFrameworkLocator.getCuratorFramework();
                    if (curator != null) {
                        try {
                            getCachedUsers(curator, path, true);
                            userInfos = users.getProperty(user);
                        } catch (Exception e) {
                            LOG.warn(e.getMessage());
                        }
                    }
                    // didn't help
                    if (userInfos == null) {
                        throw new FailedLoginException("User doesn't exist");
                    }
                }

                // the password is in the first position
                String[] infos = userInfos.split(",");
                String password = infos[0];

                if (!checkPassword(new String(tmpPassword), password)) {
                    // force reload cache of user tokens
                    CuratorFramework curator = CuratorFrameworkLocator.getCuratorFramework();
                    if (curator != null) {
                        try {
                            getCachedUsers(curator, path, true);
                            userInfos = users.getProperty(user);
                        } catch (Exception e) {
                            LOG.warn(e.getMessage());
                        }
                    }
                    // didn't help
                    if (userInfos == null) {
                        throw new FailedLoginException("User doesn't exist");
                    }
                    infos = userInfos.split(",");
                    password = infos[0];
                    if (!checkPassword(new String(tmpPassword), password)) {
                        throw new FailedLoginException("Password does not match");
                    }
                }

                principals = new HashSet<Principal>();
                principals.add(new UserPrincipal(user));
                for (int i = 1; i < infos.length; i++) {
                    if (infos[i].trim().startsWith(BackingEngine.GROUP_PREFIX)) {
                        // it's a group reference
                        principals.add(new GroupPrincipal(infos[i].trim().substring(BackingEngine.GROUP_PREFIX.length())));
                        String groupInfo = (String) users.get(infos[i].trim());
                        if (groupInfo != null) {
                            String[] roles = groupInfo.split(",");
                            for (int j = 1; j < roles.length; j++) {
                                principals.add(new RolePrincipal(roles[j].trim()));
                            }
                        }
                    } else {
                        // it's an user reference
                        principals.add(new RolePrincipal(infos[i].trim()));
                    }                }
                subject.getPrivateCredentials().add(new String(tmpPassword));
                result = true;
            }
        } catch (LoginException ex) {
            if (debug) {
                LOG.debug("Login failed {}", user, ex);
            }
            throw ex;
        }
        if (debug) {
            LOG.debug("Successfully logged in {}", user);
        }
        return result;
    }

    @Override
    public boolean commit() throws LoginException {
        if (principals.isEmpty()) {
            return false;
        }
        RolePolicy policy = RolePolicy.getPolicy(rolePolicy);
        if (policy != null && roleDiscriminator != null) {
            policy.handleRoles(subject, principals, roleDiscriminator);
        } else {
            subject.getPrincipals().addAll(principals);
        }
        return true;
    }

    public boolean abort() throws LoginException {
        if (debug) {
            LOG.debug("abort");
        }
        return true;
    }

    public boolean logout() throws LoginException {
        subject.getPrincipals().removeAll(principals);
        principals.clear();
        if (debug) {
            LOG.debug("logout");
        }
        return true;
    }

    public String getEncryptedPassword(String password) {
        Encryption encryption = encryptionSupport.getEncryption();
        String encryptionPrefix = encryptionSupport.getEncryptionPrefix();
        String encryptionSuffix = encryptionSupport.getEncryptionSuffix();

        if (encryption == null) {
            return password;
        } else {
            boolean prefix = encryptionPrefix == null || password.startsWith(encryptionPrefix);
            boolean suffix = encryptionSuffix == null || password.endsWith(encryptionSuffix);
            if (prefix && suffix) {
                return password;
            } else {
                String p = encryption.encryptPassword(password);
                if (encryptionPrefix != null) {
                    p = encryptionPrefix + p;
                }
                if (encryptionSuffix != null) {
                    p = p + encryptionSuffix;
                }
                return p;
            }
        }
    }

    public boolean checkPassword(String plain, String encrypted) {
        Encryption encryption = encryptionSupport.getEncryption();
        String encryptionPrefix = encryptionSupport.getEncryptionPrefix();
        String encryptionSuffix = encryptionSupport.getEncryptionSuffix();
        if (encryption == null) {
            return plain.equals(encrypted);
        } else {
            boolean prefix = encryptionPrefix == null || encrypted.startsWith(encryptionPrefix);
            boolean suffix = encryptionSuffix == null || encrypted.endsWith(encryptionSuffix);
            if (prefix && suffix) {
                encrypted = encrypted.substring(encryptionPrefix != null ? encryptionPrefix.length() : 0, encrypted.length()
                        - (encryptionSuffix != null ? encryptionSuffix.length() : 0));
                return encryption.checkPassword(plain, encrypted);
            } else {
                return plain.equals(encrypted);
            }
        }
    }

}
