package com.atlassian.bitbucket.internal.ssh.server;

import com.atlassian.bitbucket.auth.AuthenticationException;
import com.atlassian.bitbucket.event.auth.AuthenticationFailureEvent;
import com.atlassian.bitbucket.internal.ssh.auth.PluginSshAuthenticationHandler;
import com.atlassian.bitbucket.server.ApplicationPropertiesService;
import com.atlassian.bitbucket.ssh.util.KeyUtils;
import com.atlassian.bitbucket.user.ApplicationUser;
import com.atlassian.event.api.EventPublisher;
import com.google.common.base.Throwables;
import org.apache.mina.util.NamePreservingRunnable;
import org.apache.sshd.server.auth.pubkey.PublickeyAuthenticator;
import org.apache.sshd.server.session.ServerSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.net.SocketAddress;
import java.security.PublicKey;
import java.util.concurrent.*;

import static com.atlassian.bitbucket.internal.ssh.server.SessionAttributes.ATTRIBUTE_AUTH;
import static com.atlassian.bitbucket.scm.AuthenticationState.NOT_AUTHENTICATED;

public class DefaultPublicKeyAuthenticator implements PublickeyAuthenticator {

    public static final String AUTHENTICATION_METHOD = "ssh";
    private static final Logger log = LoggerFactory.getLogger(DefaultPublicKeyAuthenticator.class);

    private final EventPublisher eventPublisher;
    private final PluginSshAuthenticationHandler authenticationHandler;
    private final ExecutorService executorService;
    private final ApplicationPropertiesService propertiesService;

    public DefaultPublicKeyAuthenticator(EventPublisher eventPublisher, PluginSshAuthenticationHandler authenticationHandler,
                                         ExecutorService executorService, ApplicationPropertiesService propertiesService) {
        this.eventPublisher = eventPublisher;
        this.authenticationHandler = authenticationHandler;
        this.executorService = executorService;
        this.propertiesService = propertiesService;
    }

    @Override
    public boolean authenticate(final String username, final PublicKey key, final ServerSession session) {
        //Note: this method is always called in the Mina SSHD authentication thread pool. You must make sure
        //all work done here is as fast as possible and, ideally, time-boxed. Long running actions should
        //be run as part of the the ScmRequestRunnable created within SshScmRequestCommandAdapter for authentication
        //success or as a Runnable on the executorService for authentication failure

        final SocketAddress remoteAddress = session.getIoSession().getRemoteAddress();
        final SshCredentials credentials = new SshCredentials(username, key);

        Future<SshAuthentication> futureAuthResult = executorService.submit(() -> {
            Thread thread = Thread.currentThread();
            String oldName = thread.getName();

            thread.setName("ssh-auth");
            try {
                ApplicationUser user = authenticationHandler.authenticate(new DefaultSshAuthenticationContext(credentials));
                if (user == null) {
                    return null;
                }

                return new SshAuthentication(user, credentials, SshKeyDetails.getFromMDC());
            } finally {
                thread.setName(oldName);
            }
        });

        SshAuthentication authResult;
        try {
            authResult = futureAuthResult.get(propertiesService.getPluginProperty("plugin.ssh.auth.timeout", 30),
                    TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            log.info("Interrupted while authenticating SSH user ({}:{}) at {}", username,
                    KeyUtils.calculateFingerprint(key), remoteAddress);
            futureAuthResult.cancel(true);

            return false;
        } catch (TimeoutException e) {
            log.warn("Timed out while authenticating SSH user ({}:{}) at {}", username,
                    KeyUtils.calculateFingerprint(key), remoteAddress, e);
            futureAuthResult.cancel(true);

            return false;
        } catch (ExecutionException e) {
            Throwable cause = e.getCause();
            if (cause instanceof AuthenticationException) {
                // Authentication details are invalid
                log.debug("SSH user ({}:{}) at {} could not be authenticated",
                        username, KeyUtils.calculateFingerprint(key), remoteAddress,
                        log.isTraceEnabled() ? cause : null);
                handleAuthFailure(credentials, (AuthenticationException) cause);

                return false;
            }

            log.error("Error authenticating SSH user ({}:{}) at {}", username,
                    KeyUtils.calculateFingerprint(key), remoteAddress, cause);

            throw Throwables.propagate(cause);
        }

        session.setAttribute(ATTRIBUTE_AUTH, authResult);

        return true;
    }

    private void handleAuthFailure(@Nonnull final SshCredentials credentials,
                                   @Nonnull final AuthenticationException exception) {
        eventPublisher.publish(new AuthenticationFailureEvent(this,
                credentials.getUsername(), AUTHENTICATION_METHOD, exception));

        if (authenticationHandler.hasFailureHandlers()) {
            //Do this on the operation thread so we don't tie up the SSHD auth or NIO threads
            executorService.execute(new NamePreservingRunnable(() -> {
                //We don't check licensed users in our own auth handler so we should never receive a
                //NoAccessAuthenticationException and even if we did, the isInvalid hack for whoami makes it impossible
                //to handle at the SSH authentication layer, so just ignore the actual exception type
                //and hard-code it as NOT_AUTHENTICATED
                authenticationHandler.onAuthenticationFailure(new DefaultSshAuthenticationFailureContext(
                        credentials, NOT_AUTHENTICATED, exception));

            }, "ssh-scm-failure-handler"));
        }
    }
}
