SSHDaemonService.java

package automately.core.services.ssh;

import automately.core.data.User;
import automately.core.data.UserData;
import automately.core.file.VirtualFileSystem;
import automately.core.services.core.AutomatelyService;
import com.hazelcast.core.IMap;
import com.hazelcast.query.Predicates;
import io.jsync.app.core.Cluster;
import io.jsync.app.core.Logger;
import io.jsync.json.JsonObject;
import org.apache.sshd.common.NamedFactory;
import org.apache.sshd.common.compression.BuiltinCompressions;
import org.apache.sshd.common.compression.Compression;
import org.apache.sshd.server.Command;
import org.apache.sshd.server.ServerBuilder;
import org.apache.sshd.server.SshServer;
import org.apache.sshd.server.auth.pubkey.PublickeyAuthenticator;
import org.apache.sshd.server.config.keys.DefaultAuthorizedKeysAuthenticator;
import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;
import org.apache.sshd.server.session.ServerSession;

import java.io.IOException;
import java.nio.file.FileSystem;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.KeyPair;
import java.security.PublicKey;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * The SSHDaemonService is an {@link automately.core.services.core.AutomatelyService} that allows the cluster
 * to have SSH operations. This includes direct access to the {@link automately.core.file.VirtualFileSystem} using
 * SFTP.
 *
 */
public class SSHDaemonService extends AutomatelyService {

    private Cluster cluster;
    private Logger logger;
    private SshServer server = null;
    private Map<String, FileSystem> cachedFileSystems = new HashMap<>();

    @Override
    public void start(Cluster owner) {
        this.cluster = owner;
        this.logger = cluster.logger();

        JsonObject sshdConfig = coreConfig().getObject("sshd", new JsonObject());

        if(!sshdConfig.containsField("sftp_enabled")
                || !sshdConfig.containsField("port")
                || !sshdConfig.containsField("host")
                || !sshdConfig.containsField("host_key")){

            logger.info("Creating default configuration.");

            sshdConfig.putBoolean("sftp_enabled", true);
            sshdConfig.putNumber("port", 2282);
            sshdConfig.putString("host", "0.0.0.0");
            sshdConfig.putString("host_key", "sshd.ser");

            coreConfig().putObject("sshd", sshdConfig);
            cluster.config().save();
        }

        if(sshdConfig.getBoolean("sftp_enabled", false)){

            logger.info("The SSH Daemon is enabled. This is experimental right now.");

            server = ServerBuilder.builder().build(true);

            server.setFileSystemFactory(session -> {
                try {
                    String username = session.getUsername();
                    if(cachedFileSystems.containsKey(username)){
                        return cachedFileSystems.get(username);
                    }
                    Collection<User> users = users().values(Predicates.equal("username", username));
                    if(users.size() == 0){
                        throw new IOException("The user \"" + session.getUsername() + "\" could not be found.");
                    }
                    FileSystem fs = VirtualFileSystem.getUserFileSystem(users.iterator().next());
                    cachedFileSystems.put(username, fs);
                    return fs;
                } catch (Exception e){
                    e.printStackTrace();
                }
                return null;
            });

            server.setCommandFactory(new SSHCommandFactory());

            List<NamedFactory<Command>> subSystemFactoryList = new ArrayList<>();
            subSystemFactoryList.add(new SubSystemFactory());

            server.setSubsystemFactories(subSystemFactoryList);

            List<NamedFactory<Compression>> compressionFactoryList = new ArrayList<>();
            compressionFactoryList.add(BuiltinCompressions.none);
            compressionFactoryList.add(BuiltinCompressions.zlib);
            compressionFactoryList.add(BuiltinCompressions.delayedZlib);

            server.setCompressionFactories(compressionFactoryList);

            IMap<String, Integer> authAttempts = cluster().data().getMap("ssh.auth.attempts");

            // TODO Finish implementing brute force protection
            server.setPublickeyAuthenticator(new PublickeyAuthenticator() {
                private Map<String, DefaultAuthorizedKeysAuthenticator> authenticatorMap = new ConcurrentHashMap<>();
                @Override
                public boolean authenticate(String username, PublicKey publicKey, ServerSession serverSession) {
                    if(!authenticatorMap.containsKey(username)){
                        try {
                            serverSession.setUsername(username);
                            FileSystem fs = server.getFileSystemFactory().createFileSystem(serverSession);
                            Path authorizedFiles = fs.getPath("/private/authorized_keys");
                            authenticatorMap.put(username, new DefaultAuthorizedKeysAuthenticator(username, authorizedFiles, false));
                        } catch (IOException e) {
                            return false;
                        }
                    }
                    return authenticatorMap.get(username).authenticate(username, publicKey, serverSession);
                }
            });

            server.setPasswordAuthenticator((user, passwordOrKey, serverSession) -> {

                User mUser = UserData.getUserByUsername(user);

                // If there has been more than 10 failed attempts within the next 3 minutes we need to deny it..
                if (mUser != null && authAttempts.getOrDefault(mUser.username, 0) < 10) {
                    logger.info("Password Authentication: " + mUser.username);
                    boolean passed = UserData.validateUserPassword(mUser, passwordOrKey) ||
                            UserData.validateUserKey(mUser, passwordOrKey) || cluster.config().isDebug();

                    if (!passed) {
                        logger.info("Authentication Failed: " + mUser.username);
                        Integer attempts = authAttempts.getOrDefault(mUser.username, 0);
                        attempts++;
                        authAttempts.put(mUser.username, attempts, 3, TimeUnit.MINUTES);
                    } else {
                        logger.info("Authentication Success: " + mUser.username);
                    }
                    return passed;
                }
                return false;
            });

            // By default it is set to listen on localhost for secure purposes.
            server.setHost(sshdConfig.getString("host", "127.0.0.1"));
            server.setPort(sshdConfig.getNumber("port", 2282).intValue());

            String hostKey = sshdConfig.getString("host_key","sshd_hostkey.ser");

            logger.info("Using the hostkey " + hostKey);

            SimpleGeneratorHostKeyProvider generator = new SimpleGeneratorHostKeyProvider(Paths.get(hostKey));
            generator.setKeySize(4098);
            generator.setAlgorithm("RSA");

            List<KeyPair> loadedKeys = generator.loadKeys();

            if(loadedKeys.size() > 0){
                KeyPair keyPair = loadedKeys.get(0);
                byte[] encodedPublicKey = keyPair.getPublic().getEncoded();
                JsonObject data = new JsonObject();
                data.putString("algorithm", keyPair.getPublic().getAlgorithm());
                data.putBinary("bytes", encodedPublicKey);

                cluster().data().persistentMap("sshd.service.storage").set("public_key", data);
            }

            server.setKeyPairProvider(generator);
            try {
                logger.info("Starting the SshServer Daemon on port " + server.getPort());
                server.start();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public void stop() {
        if(server != null){
            try {
                server.stop(true);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public String name() {
        return getClass().getCanonicalName();
    }
}