package ru.linachan.yggdrasil.shell;


import org.apache.sshd.server.forward.RejectAllForwardingFilter;
import ru.linachan.yggdrasil.auth.YggdrasilAuthUser;
import ru.linachan.yggdrasil.service.YggdrasilService;

import org.apache.sshd.server.*;
import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;

import java.io.File;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

public class YggdrasilShellService extends YggdrasilService {

    private SshServer shellServer;

    @Override
    @SuppressWarnings("unchecked")
    protected void onInit() {
        YggdrasilShellCommandManager commandManager = new YggdrasilShellCommandManager();
        commandManager.setUpManager(core);

        shellServer = SshServer.setUpDefaultServer();

        shellServer.setPort(core.getConfig().getInt("yggdrasil.ssh.port", 41598));
        shellServer.setHost(core.getConfig().getString("yggdrasil.ssh.host", "0.0.0.0"));

        shellServer.getProperties().put(SshServer.IDLE_TIMEOUT, 86400000L);

        shellServer.setKeyPairProvider(new SimpleGeneratorHostKeyProvider(
            new File(core.getConfig().getString("yggdrasil.ssh.key", "master.ser"))
        ));

        shellServer.setPublickeyAuthenticator((userName, publicKey, serverSession) -> {
            YggdrasilAuthUser authUser = core.getAuthManager().getUser(userName);

            if (authUser != null) {
                String publicKeyString = new String(Base64.getEncoder().encode(publicKey.getEncoded()));

                for (String publicKeyData: ((List<String>) authUser.getAttribute("publicKey"))) {
                    if (publicKeyString.equals(publicKeyData)) {
                        return true;
                    }
                }

                authUser.setAttribute("tmpPublicKey", publicKeyString);
            } else if (core.getConfig().getString("yggdrasil.admin.user", "yggdrasil").equals(userName)) {
                authUser = core.getAuthManager().registerUser(userName);

                List<String> publicKeyList = new ArrayList<>();
                publicKeyList.add(new String(Base64.getEncoder().encode(publicKey.getEncoded())));
                authUser.setAttribute("publicKey", publicKeyList);

                core.getAuthManager().updateUser(authUser);

                return true;
            }

            return false;
        });

        shellServer.setPasswordAuthenticator((userName, passWord, serverSession) -> {
            YggdrasilAuthUser authUser = core.getAuthManager().getUser(userName);

            if (authUser != null) {
                MessageDigest md = null;
                try {
                    md = MessageDigest.getInstance("SHA-256");
                    md.update(passWord.getBytes("UTF-8"));
                    byte[] digest = md.digest();
                    return String.format("%064x", new java.math.BigInteger(1, digest))
                        .equals(authUser.getAttribute("passWord"));
                } catch (UnsupportedEncodingException | NoSuchAlgorithmException e) {
                    logger.error("Unable to calculate SHA-256: {}", e.getMessage());
                }

                return false;
            } else if (core.getConfig().getString("yggdrasil.admin.user", "yggdrasil").equals(userName)) {
                authUser = core.getAuthManager().registerUser(userName);
                authUser.setAttribute("passWord", passWord);
                core.getAuthManager().updateUser(authUser);

                return true;
            }

            return false;
        });

        shellServer.setTcpipForwardingFilter(RejectAllForwardingFilter.INSTANCE);

        shellServer.setCommandFactory(new YggdrasilShellCommandFactory(core, commandManager));
        shellServer.setShellFactory(new YggdrasilShellFactory(core, commandManager));

        try {
            shellServer.start();
        } catch (IOException e) {
            logger.error("Unable to start MasterPlugin SSH server", e);
        }
    }

    @Override
    protected void onShutdown() {
        try {
            shellServer.stop();
        } catch (IOException e) {
            logger.error("Unable to stop MasterPlugin SSH server properly", e);
        }
    }

    @Override
    public void run() {
        // do nothing, service already working
    }
}

@@ -93,24 +93,100 @@
 dingFilter;%0A
+import org.apache.sshd.server.keyprovider.AbstractGeneratorHostKeyProvider;%0A
 import ru.li
@@ -1274,38 +1274,56 @@
     
-shellServer.setKeyPair
+AbstractGeneratorHostKeyProvider hostKey
 Provider
 (new
@@ -1318,17 +1318,19 @@
 Provider
-(
+ = 
 new Simp
@@ -1450,16 +1450,119 @@
        )
+;%0A%0A        hostKeyProvider.setAlgorithm(%22RSA%22);%0A%0A        shellServer.setKeyPairProvider(hostKeyProvider
 );%0A%0A    

Make YggdrasilShell use RSA instead of DSS