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

import com.atlassian.bitbucket.internal.ssh.service.SshConfigurationRefreshEvent;
import com.atlassian.bitbucket.server.ApplicationPropertiesService;
import com.atlassian.bitbucket.ssh.SshConfiguration;
import com.atlassian.bitbucket.ssh.SshConfigurationService;
import com.atlassian.bitbucket.ssh.event.SshConfigurationChangedEvent;
import com.atlassian.event.api.EventListener;
import com.atlassian.sal.api.lifecycle.LifecycleAware;
import com.atlassian.util.concurrent.ThreadFactories;
import com.google.common.collect.ImmutableList;
import org.apache.sshd.common.compression.BuiltinCompressions;
import org.apache.sshd.common.io.mina.MinaServiceFactoryFactory;
import org.apache.sshd.common.session.Session;
import org.apache.sshd.common.session.SessionListener;
import org.apache.sshd.common.util.SecurityUtils;
import org.apache.sshd.server.CommandFactory;
import org.apache.sshd.server.ServerBuilder;
import org.apache.sshd.server.ServerFactoryManager;
import org.apache.sshd.server.auth.pubkey.PublickeyAuthenticator;
import org.apache.sshd.server.global.KeepAliveHandler;
import org.apache.sshd.server.global.NoMoreSessionsHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.core.Ordered;

import javax.annotation.Nonnull;
import javax.management.JMException;
import javax.management.ObjectName;
import java.io.IOException;
import java.io.InterruptedIOException;
import java.lang.management.ManagementFactory;
import java.util.Arrays;
import java.util.Optional;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import static com.google.common.base.Preconditions.checkState;
import static java.util.Optional.empty;
import static java.util.Optional.of;

/**
 * Wraps and manages the Mina/SSHD server.
 */
public class SshServer implements DisposableBean, LifecycleAware, Ordered {

    private static final Logger log = LoggerFactory.getLogger(SshServer.class);
    private static final long DEFAULT_IDLE_TIMEOUT = TimeUnit.HOURS.toSeconds(2);
    private static final String MXBEAN_NAME = "com.atlassian.bitbucket:name=SshSessions";
    private static final String PROP_IDLE_TIMEOUT = "plugin.ssh.command.timeout.idle";
    private static final String PROP_SSH_NIO_WORKERS = "plugin.ssh.nio.workers";

    private final boolean jmxEnabled;
    private final Object lock;
    private final org.apache.sshd.server.SshServer sshServer;
    private final SshConfigurationService sshConfigurationService;

    private volatile SessionTracker sessionTracker;
    private volatile SshServerStatus status;

    public SshServer(PublickeyAuthenticator authenticator, CommandFactory commandFactory,
                     SshConfigurationService configurationService, HostKeyPairProvider keyPairProvider,
                     ApplicationPropertiesService propertiesService, NamedFactoryFilter factoryFilter)
            throws IOException {
        // Without BouncyCastle SSHD may not behave as expected - see BSERVDEV-6984
        checkState(SecurityUtils.isBouncyCastleRegistered(), "BouncyCastle is not registered as expected");

        this.sshConfigurationService = configurationService;

        jmxEnabled = propertiesService.isJmxEnabled();
        lock = new Object();
        sshServer = init(authenticator, commandFactory, keyPairProvider, propertiesService, factoryFilter);
        status = SshServerStatus.forState(SshServerState.STOPPED);
    }

    @Override
    public void destroy() {
        stopServer();
    }

    @Override
    public int getOrder() {
        return 10; // Plugin upgrades run <= 5. SSH should not start until upgrades are complete
    }

    /**
     * @return the current {@link SshServerStatus status} of the SSH server
     * @since 3.7
     */
    public SshServerStatus getStatus() {
        return status;
    }

    /**
     * @return {@code true} if the SSH server is running; otherwise, {@code false} if it has been stopped or
     *         failed to start
     * @since 3.7
     */
    public boolean isRunning() {
        return status.getState() == SshServerState.RUNNING;
    }

    @EventListener
    public void onConfigurationChange(SshConfigurationChangedEvent event) {
        synchronizeStateFrom(event.getNewConfiguration());
    }

    @EventListener
    public void onConfigurationRefresh(SshConfigurationRefreshEvent event) {
        synchronizeStateFrom(sshConfigurationService.getConfiguration());
    }

    /**
     * When {@code LifecycleAware} indicates the application has started, start the SSH server if it is enabled.
     * <p>
     * Note: This should <i>not</i> be done in response to the {@code ApplicationContext} (ruling out both
     * {@code InitializingBean} and {@code ApplicationListener}) because SSH should not start until plugins
     * have been upgraded (which happens in response to the application's lifecycle).
     */
    @Override
    public void onStart() {
        if (jmxEnabled) {
            registerMxBean();
        }
        SshConfiguration sshConfiguration = sshConfigurationService.getConfiguration();
        if (sshConfiguration.isEnabled()) {
            startServer(sshConfiguration.getPort());
        }
    }

    @Override
    public void onStop() {
        if (jmxEnabled) {
            unregisterMxBean();
        }
    }

    public void startServer(int port) {
        if (status.getState() != SshServerState.RUNNING) {
            synchronized (lock) {
                if (status.getState() != SshServerState.RUNNING) {
                    startServerInternal(port);
                }
            }
        }
    }

    public void stopServer() {
        if (status.getState() != SshServerState.STOPPED) {
            synchronized (lock) {
                if (status.getState() != SshServerState.STOPPED) {
                    stopServerInternal();
                }
            }
        }
    }

    public Optional<Integer> getPort() {
        return status.getState() == SshServerState.RUNNING ? of(sshServer.getPort()) : empty();
    }

    @Nonnull
    private static ScheduledExecutorService createScheduledExecutor(String name) {
        // This executor is identical to the default one (see org.apache.sshd.SshServer.checkConfig), except
        // that this will create daemon threads where the default one will not
        return Executors.newSingleThreadScheduledExecutor(
                ThreadFactories.namedThreadFactory(name, ThreadFactories.Type.DAEMON));
    }

    private org.apache.sshd.server.SshServer init(PublickeyAuthenticator authenticator,
                                                  CommandFactory scmCommandFactory,
                                                  HostKeyPairProvider keyPairProvider,
                                                  ApplicationPropertiesService propertiesService,
                                                  NamedFactoryFilter factoryFilter) {
        org.apache.sshd.server.SshServer server = new ServerBuilder()
                // Prevent filesystem access
                .fileSystemFactory(new ForbiddingFileSystemFactory())
                // Prevent TCP/IP forwarding handlers from being registered (they're two of the default handlers)
                .globalRequestHandlers(Arrays.asList(new KeepAliveHandler(), new NoMoreSessionsHandler()))
                .build();

        server.setChannelFactories(ImmutableList.of(
                // Ensure channels are closed, since SSHD has some gaps
                new SelfClosingChannelSessionFactory(TimeUnit.SECONDS.toMillis(
                        propertiesService.getPluginProperty(PROP_IDLE_TIMEOUT, DEFAULT_IDLE_TIMEOUT)))
        ));
        server.setCipherFactories(factoryFilter.filter(server.getCipherFactories(), NamedFactoryType.CIPHER));
        server.setCommandFactory(scmCommandFactory);
        server.setCompressionFactories(ImmutableList.of(BuiltinCompressions.none));
        server.setGSSAuthenticator(null); // Authentication should require a public key
        server.setIoServiceFactoryFactory(new MinaServiceFactoryFactory()); // BSERVDEV-7214 Use Mina, not NIO
        server.setKeyExchangeFactories(factoryFilter.filter(server.getKeyExchangeFactories(), NamedFactoryType.KEY_EXCHANGE));
        server.setKeyPairProvider(keyPairProvider);
        server.setKeyboardInteractiveAuthenticator(null); // Authentication should require a public key
        server.setMacFactories(factoryFilter.filter(server.getMacFactories(), NamedFactoryType.MAC));
        server.setNioWorkers(propertiesService.getPluginProperty(PROP_SSH_NIO_WORKERS, ServerFactoryManager.DEFAULT_NIO_WORKERS));
        server.setPasswordAuthenticator(null); // Authentication should require a public key
        server.setPublickeyAuthenticator(authenticator);
        server.setScheduledExecutorService(createScheduledExecutor("ssh-server"), true);
        server.setShellFactory(null); // Remote users should not be able to get an interactive shell
        server.setTcpipForwardingFilter(null); // Remote users should not be able to request TCP/IP forwarding

        return server;
    }

    private void registerMxBean() {
        if (sessionTracker == null) {
            sessionTracker = new SessionTracker();
        }
        sshServer.addSessionListener(sessionTracker);

        try {
            ManagementFactory.getPlatformMBeanServer().registerMBean(
                    new SshSessionsMxBeanAdapter(sessionTracker), new ObjectName(MXBEAN_NAME));
        } catch (RuntimeException | JMException e) {
            log.warn("Could not registers SshSessions MXBean. SSH details will not be available in JMX", e);

            sshServer.removeSessionListener(sessionTracker);
            sessionTracker = null;
        }
    }

    private void startServerInternal(int port) {
        try {
            log.info("Starting SSH server on port {}...", port);
            sshServer.setPort(port);
            sshServer.start();
            status = SshServerStatus.forState(SshServerState.RUNNING);
            log.info("Started SSH server successfully.");
        } catch (IOException e) {
            status = SshServerStatus.forException(e);
            log.error("SSH server failed to start", e);
        }
    }

    private void stopServerInternal() {
        if (isRunning()) {
            try {
                log.info("Stopping SSH server...");
                sshServer.stop(true);
                log.info("Stopped SSH server successfully.");
            } catch (InterruptedIOException e) {
                // This shouldn't happen, since we've told SshServer to stop immediately
                log.warn("Interrupted while waiting for SSH server to stop; some sessions may still be active", e);
            } catch (IOException e) {
                log.warn("The SSH server could not be stopped", e);
            }
        }
        // This ensures we clear the FAILED_TO_START flag if explicitly stopped
        status = SshServerStatus.forState(SshServerState.STOPPED);
    }

    private void synchronizeStateFrom(SshConfiguration configuration) {
        if (isRunning()) {
            if (!configuration.isEnabled()) {
                stopServer();
            } else if (sshServer.getPort() != configuration.getPort()) {
                stopServer();
                startServer(configuration.getPort());
            }
        } else if (configuration.isEnabled()) {
            startServer(configuration.getPort());
        }
    }

    private void unregisterMxBean() {
        sshServer.removeSessionListener(sessionTracker);

        try {
            ManagementFactory.getPlatformMBeanServer().unregisterMBean(new ObjectName(MXBEAN_NAME));
        } catch (RuntimeException | JMException e) {
            log.warn("Failed to unregister SshSessions MXBean", e);
        }

    }

    private static class SessionTracker implements SessionListener {

        private final AtomicLong maxActiveSessionCount;
        private final AtomicLong sessionClosedCount;
        private final AtomicLong sessionCreatedCount;
        private final AtomicLong sessionExceptionCount;

        SessionTracker() {
            maxActiveSessionCount = new AtomicLong();
            sessionClosedCount = new AtomicLong();
            sessionCreatedCount = new AtomicLong();
            sessionExceptionCount = new AtomicLong();
        }

        long getActiveSessionCount() {
            return sessionCreatedCount.get() - sessionClosedCount.get();
        }

        long getMaxActiveSessionCount() {
            return maxActiveSessionCount.get();
        }

        long getSessionClosedCount() {
            return sessionClosedCount.get();
        }

        long getSessionCreatedCount() {
            return sessionCreatedCount.get();
        }

        long getSessionExceptionCount() {
            return sessionExceptionCount.get();
        }

        @Override
        public void sessionCreated(Session session) {
            long created = sessionCreatedCount.incrementAndGet();
            long active = created - sessionClosedCount.get();
            maxActiveSessionCount.accumulateAndGet(active, Math::max);
        }

        @Override
        public void sessionEvent(Session session, Event event) {
            // no-op
        }

        @Override
        public void sessionException(Session session, Throwable t) {
            sessionExceptionCount.incrementAndGet();
        }

        @Override
        public void sessionClosed(Session session) {
            sessionClosedCount.incrementAndGet();
        }
    }

    private static class SshSessionsMxBeanAdapter implements SshSessionsMXBean {

        private final SessionTracker sessionTracker;

        private SshSessionsMxBeanAdapter(SessionTracker sessionTracker) {
            this.sessionTracker = sessionTracker;
        }

        public long getActiveSessionCount() {
            return sessionTracker.getActiveSessionCount();
        }

        public long getMaxActiveSessionCount() {
            return sessionTracker.getMaxActiveSessionCount();
        }

        public long getSessionClosedCount() {
            return sessionTracker.getSessionClosedCount();
        }

        public long getSessionCreatedCount() {
            return sessionTracker.getSessionCreatedCount();
        }

        public long getSessionExceptionCount() {
            return sessionTracker.getSessionExceptionCount();
        }

    }
}
