package com.atlassian.bitbucket.mesh.test.util.ssh;

import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.repository.RepositoryManager;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import org.apache.commons.lang3.StringUtils;
import org.apache.sshd.common.compression.BuiltinCompressions;
import org.apache.sshd.common.util.net.SshdSocketAddress;
import org.apache.sshd.server.ServerBuilder;
import org.apache.sshd.server.SshServer;
import org.apache.sshd.server.auth.AsyncAuthException;
import org.apache.sshd.server.channel.ChannelSession;
import org.apache.sshd.server.command.Command;
import org.apache.sshd.server.keyprovider.SimpleGeneratorHostKeyProvider;
import org.apache.sshd.server.session.ServerSession;
import org.bouncycastle.openssl.jcajce.JcaPEMWriter;

import javax.annotation.Nonnull;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.util.*;

import static java.util.Objects.requireNonNull;

/**
 * A minimal SSH server that can be used to test git commands that access a remote server, such as {@code git fetch}
 * and {@code git push}. The server only supports 2 commands {@code git-upload-pack} and {@code git-receive-pack}.
 */
public class TestGitSshServer implements AutoCloseable {

    private final GitProcessBuilderFactory builderFactory;
    private final RepositoryManager repositoryManager;
    private final Map<String, PublicKey> userKeys;

    private SshServer sshServer;

    public TestGitSshServer(GitProcessBuilderFactory builderFactory, RepositoryManager repositoryManager) {
        this.builderFactory = builderFactory;
        this.repositoryManager = repositoryManager;

        userKeys = new HashMap<>();
    }

    /**
     * Registers the provided {@code publicKey} for the {@code username}, which allows the user to authenticate with
     * the associated private key.
     *
     * @param username the username
     * @param publicKey the public key
     */
    public void addCredentials(@Nonnull String username, @Nonnull PublicKey publicKey) {
        userKeys.put(requireNonNull(username, "username"), requireNonNull(publicKey, "publicKey"));
    }

    /**
     * Generates an SSH key pair for {@code username}, registers the public key with the SSH authenticator and
     * returns the PEM private key as a {@code ByteString}.
     *
     * @param username the username
     * @return the PEM private key as a {@code ByteString}
     */
    public ByteString generateUserCredentials(@Nonnull String username) {
        KeyPair keyPair = generateKeyPair();
        addCredentials(username, keyPair.getPublic());
        return getPrivateKeyPemBytes(keyPair.getPrivate());
    }

    @Override
    public void close() throws IOException {
        stop();
    }

    /**
     * Generates an SSH clone URL that can be used to access the repository over SSH. To avoid unsafe characters in the
     * clone URL the {@code repositoryID} string is base64 encoded. The resulting clone URL is of the form
     * {@code ssh://[username] [at] localhost:[ssh-port]/[base64-encoded-repository-id]}. Pardon the use of [at] here,
     * it's written in this way to avoid the security scanner tripping over a potentially unsafe URL in the code.
     *
     * @param username the username
     * @param repositoryId the ID of the repository for which to generate a clone URL
     * @return the clone URL for the provided user and repository
     */
    public String getCloneUrl(@Nonnull String username, @Nonnull String repositoryId) {
        String path = Base64.getEncoder().encodeToString(
                requireNonNull(repositoryId, "repositoryId").getBytes(StandardCharsets.UTF_8));
        return "ssh://" + requireNonNull(username, "username") + "@localhost:" + sshServer.getPort() + "/" + path;
    }

    public OptionalInt getSshPort() {
        return sshServer != null ? OptionalInt.of(sshServer.getPort()) : OptionalInt.empty();
    }

    public void reset() {
        userKeys.clear();
    }

    public void start() throws IOException {
        if (sshServer == null) {
            sshServer = new ServerBuilder()
                    .publickeyAuthenticator(this::authenticateKey)
                    .compressionFactories(ImmutableList.of(BuiltinCompressions.none))
                    .build();
            sshServer.setCommandFactory(this::createSshCommand);
            sshServer.setKeyPairProvider(new SimpleGeneratorHostKeyProvider());
            sshServer.setHost(SshdSocketAddress.LOCALHOST_NAME);
            sshServer.start();
        }
    }

    public void stop() throws IOException {
        if (sshServer != null) {
            sshServer.stop();
            sshServer = null;
        }
    }

    private static KeyPair generateKeyPair() {
        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
            keyGen.initialize(2048);
            return keyGen.generateKeyPair();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    private static ByteString getPrivateKeyPemBytes(PrivateKey privateKey) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream(4096);
        try (Writer writer = new OutputStreamWriter(outputStream, StandardCharsets.US_ASCII);
             JcaPEMWriter pemWriter = new JcaPEMWriter(writer)) {
            pemWriter.writeObject(privateKey);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
        return ByteString.copyFrom(outputStream.toByteArray());
    }

    private boolean authenticateKey(String username, PublicKey key, ServerSession session) throws AsyncAuthException {
        return Objects.equals(key, userKeys.get(username));
    }

    private Command createSshCommand(ChannelSession channel, String command) {
        //command should look like this: git-upload-pack '/<base64-encoded-repository-id>'
        if (!command.startsWith("git-upload-pack ") && !command.startsWith("git-receive-pack ")) {
            // unsupported command
            return null;
        }

        int index = command.indexOf(' ');
        //Strip off quotes and any leading slashes
        String arguments = StringUtils.strip(command.substring(index + 1), "\\'\"/");
        Repository repository = getRepository(arguments);
        //Use a builder not tied to a repository so we don't get auto-config magic
        return new GitSshCommand(builderFactory.builder()
                .workDir(repository.getPath())
                .command(command.substring(4, index)) // drop the git- prefix to get just upload-pack or receive-pack
                .argument(repository.getPath().toString())
        );
    }

    private Repository getRepository(String arguments) {
        String decodedId = new String(Base64.getDecoder().decode(StringUtils.trim(arguments)), StandardCharsets.UTF_8);
        return repositoryManager.getById(decodedId);
    }
}
