code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.base.ssh;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import com.google.common.base.Preconditions;
  import io.airbyte.commons.functional.CheckedConsumer;
  import io.airbyte.commons.functional.CheckedFunction;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.commons.string.Strings;
  import java.io.IOException;
  import java.io.StringReader;
  import java.net.InetSocketAddress;
  import java.security.KeyPair;
  import java.util.List;
  import org.apache.sshd.client.SshClient;
  import org.apache.sshd.client.keyverifier.AcceptAllServerKeyVerifier;
  import org.apache.sshd.client.session.ClientSession;
  import org.apache.sshd.common.util.net.SshdSocketAddress;
  import org.apache.sshd.server.forward.AcceptAllForwardingFilter;
  import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
  import org.bouncycastle.openssl.PEMKeyPair;
  import org.bouncycastle.openssl.PEMParser;
  import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  // todo (cgardens) - this needs unit tests. it is currently tested transitively via source postgres
  // integration tests.
  /**
   * Encapsulates the connection configuration for an ssh tunnel port forward through a proxy/bastion
   * host plus the remote host and remote port to forward to a specified local port.
   */
  public class SshTunnel implements AutoCloseable {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(SshTunnel.class);
  
    public enum TunnelMethod {
      NO_TUNNEL,
      SSH_PASSWORD_AUTH,
      SSH_KEY_AUTH
    }
  
    public static final int TIMEOUT_MILLIS = 15000; // 15 seconds
  
    private final JsonNode config;
    private final List<String> hostKey;
    private final List<String> portKey;
  
    private final TunnelMethod tunnelMethod;
    private final String tunnelHost;
    private final int tunnelPort;
    private final String tunnelUser;
    private final String sshKey;
    private final String tunnelUserPassword;
    private final String remoteDatabaseHost;
    private final int remoteDatabasePort;
    private int tunnelDatabasePort;
  
    private SshClient sshclient;
    private ClientSession tunnelSession;
  
    /**
     *
     * @param config - the full config that was passed to the source.
     * @param hostKey - a list of keys that point to the database host name. should be pointing to where
     *        in the config remoteDatabaseHost is found.
     * @param portKey - a list of keys that point to the database port. should be pointing to where in
     *        the config remoteDatabasePort is found.
     * @param tunnelMethod - the type of ssh method that should be used (includes not using SSH at all).
     * @param tunnelHost - host name of the machine to which we will establish an ssh connection (e.g.
     *        hostname of the bastion).
     * @param tunnelPort - port of the machine to which we will establish an ssh connection. (e.g. port
     *        of the bastion).
     * @param tunnelUser - user that is allowed to access the tunnelHost.
     * @param sshKey - the ssh key that will be used to make the ssh connection. can be null if we are
     *        using tunnelUserPassword instead.
     * @param tunnelUserPassword - the password for the tunnelUser. can be null if we are using sshKey
     *        instead.
     * @param remoteDatabaseHost - the actual host name of the database (as it is known to the tunnel
     *        host).
     * @param remoteDatabasePort - the actual port of the database (as it is known to the tunnel host).
     */
    public SshTunnel(final JsonNode config,
                     final List<String> hostKey,
                     final List<String> portKey,
                     final TunnelMethod tunnelMethod,
                     final String tunnelHost,
                     final int tunnelPort,
                     final String tunnelUser,
                     final String sshKey,
                     final String tunnelUserPassword,
                     final String remoteDatabaseHost,
                     final int remoteDatabasePort) {
      this.config = config;
      this.hostKey = hostKey;
      this.portKey = portKey;
  
      Preconditions.checkNotNull(tunnelMethod);
      this.tunnelMethod = tunnelMethod;
  
      if (tunnelMethod.equals(TunnelMethod.NO_TUNNEL)) {
        this.tunnelHost = null;
        this.tunnelPort = 0;
        this.tunnelUser = null;
        this.sshKey = null;
        this.tunnelUserPassword = null;
        this.remoteDatabaseHost = null;
        this.remoteDatabasePort = 0;
      } else {
        Preconditions.checkNotNull(tunnelHost);
        Preconditions.checkArgument(tunnelPort > 0);
        Preconditions.checkNotNull(tunnelUser);
        if (tunnelMethod.equals(TunnelMethod.SSH_KEY_AUTH)) {
          Preconditions.checkNotNull(sshKey);
        }
        if (tunnelMethod.equals(TunnelMethod.SSH_PASSWORD_AUTH)) {
          Preconditions.checkNotNull(tunnelUserPassword);
        }
        Preconditions.checkNotNull(remoteDatabaseHost);
        Preconditions.checkArgument(remoteDatabasePort > 0);
  
        this.tunnelHost = tunnelHost;
        this.tunnelPort = tunnelPort;
        this.tunnelUser = tunnelUser;
        this.sshKey = sshKey;
        this.tunnelUserPassword = tunnelUserPassword;
        this.remoteDatabaseHost = remoteDatabaseHost;
        this.remoteDatabasePort = remoteDatabasePort;
  
        this.sshclient = createClient();
        this.tunnelSession = openTunnel(sshclient);
      }
    }
  
    public JsonNode getOriginalConfig() {
      return config;
    }
  
    public JsonNode getConfigInTunnel() {
      if (tunnelMethod.equals(TunnelMethod.NO_TUNNEL)) {
        return getOriginalConfig();
      } else {
        final JsonNode clone = Jsons.clone(config);
        Jsons.replaceNestedString(clone, hostKey, SshdSocketAddress.LOCALHOST_ADDRESS.getHostName());
        Jsons.replaceNestedInt(clone, portKey, tunnelDatabasePort);
        return clone;
      }
    }
  
    // /**
    // * Finds a free port on the machine. As soon as this method returns, it is possible for process to
    // bind to this port. Thus it only gives a guarantee that at the time
    // */
    // private static int findFreePort() {
    // // finds an available port.
    // try (final var socket = new ServerSocket(0)) {
    // return socket.getLocalPort();
    // } catch (final IOException e) {
    // throw new RuntimeException(e);
    // }
    // }
  
    public static SshTunnel getInstance(final JsonNode config, final List<String> hostKey, final List<String> portKey) {
      final TunnelMethod tunnelMethod = Jsons.getOptional(config, "tunnel_method", "tunnel_method")
          .map(method -> TunnelMethod.valueOf(method.asText().trim()))
          .orElse(TunnelMethod.NO_TUNNEL);
      LOGGER.info("Starting connection with method: {}", tunnelMethod);
  
      // final int localPort = findFreePort();
  
      return new SshTunnel(
          config,
          hostKey,
          portKey,
          tunnelMethod,
          Strings.safeTrim(Jsons.getStringOrNull(config, "tunnel_method", "tunnel_host")),
          Jsons.getIntOrZero(config, "tunnel_method", "tunnel_port"),
          Strings.safeTrim(Jsons.getStringOrNull(config, "tunnel_method", "tunnel_user")),
          Strings.safeTrim(Jsons.getStringOrNull(config, "tunnel_method", "ssh_key")),
          Strings.safeTrim(Jsons.getStringOrNull(config, "tunnel_method", "tunnel_user_password")),
          Strings.safeTrim(Jsons.getStringOrNull(config, hostKey)),
          Jsons.getIntOrZero(config, portKey));
    }
  
    public static void sshWrap(final JsonNode config,
                               final List<String> hostKey,
                               final List<String> portKey,
                               final CheckedConsumer<JsonNode, Exception> wrapped)
        throws Exception {
      sshWrap(config, hostKey, portKey, (configInTunnel) -> {
        wrapped.accept(configInTunnel);
        return null;
      });
    }
  
    public static <T> T sshWrap(final JsonNode config,
                                final List<String> hostKey,
                                final List<String> portKey,
                                final CheckedFunction<JsonNode, T, Exception> wrapped)
        throws Exception {
      try (final SshTunnel sshTunnel = SshTunnel.getInstance(config, hostKey, portKey)) {
        return wrapped.apply(sshTunnel.getConfigInTunnel());
      }
    }
  
    /**
     * Closes a tunnel if one was open, and otherwise doesn't do anything (safe to run).
     */
    @Override
    public void close() {
      try {
        if (tunnelSession != null) {
          tunnelSession.close();
          tunnelSession = null;
        }
        if (sshclient != null) {
          sshclient.stop();
          sshclient = null;
        }
      } catch (final Throwable t) {
        throw new RuntimeException(t);
      }
    }
  
    /**
     * From the RSA format private key string, use bouncycastle to deserialize the key pair, reconstruct
     * the keys from the key info, and return the key pair for use in authentication.
     */
    private KeyPair getPrivateKeyPair() throws IOException {
      final PEMParser pemParser = new PEMParser(new StringReader(validateKey()));
      final PEMKeyPair keypair = (PEMKeyPair) pemParser.readObject();
      final JcaPEMKeyConverter converter = new JcaPEMKeyConverter();
      return new KeyPair(
          converter.getPublicKey(SubjectPublicKeyInfo.getInstance(keypair.getPublicKeyInfo())),
          converter.getPrivateKey(keypair.getPrivateKeyInfo()));
    }
  
    private String validateKey() {
      return sshKey.replace("\\n", "\n");
    }
  
    /**
     * Generates a new ssh client and returns it, with forwarding set to accept all types; use this
     * before opening a tunnel.
     */
    private SshClient createClient() {
      java.security.Security.addProvider(
          new org.bouncycastle.jce.provider.BouncyCastleProvider());
      final SshClient client = SshClient.setUpDefaultClient();
      client.setForwardingFilter(AcceptAllForwardingFilter.INSTANCE);
      client.setServerKeyVerifier(AcceptAllServerKeyVerifier.INSTANCE);
      return client;
    }
  
    /**
     * Starts an ssh session; wrap this in a try-finally and use closeTunnel() to close it.
     */
    private ClientSession openTunnel(final SshClient client) {
      try {
        client.start();
        final ClientSession session = client.connect(
            tunnelUser.trim(),
            tunnelHost.trim(),
            tunnelPort)
            .verify(TIMEOUT_MILLIS)
            .getSession();
        if (tunnelMethod.equals(TunnelMethod.SSH_KEY_AUTH)) {
          session.addPublicKeyIdentity(getPrivateKeyPair());
        }
        if (tunnelMethod.equals(TunnelMethod.SSH_PASSWORD_AUTH)) {
          session.addPasswordIdentity(tunnelUserPassword);
        }
  
        session.auth().verify(TIMEOUT_MILLIS);
        final SshdSocketAddress address = session.startLocalPortForwarding(
            // entering 0 lets the OS pick a free port for us.
            new SshdSocketAddress(InetSocketAddress.createUnresolved(SshdSocketAddress.LOCALHOST_ADDRESS.getHostName(), 0)),
            new SshdSocketAddress(remoteDatabaseHost, remoteDatabasePort));
  
        // discover the port that the OS picked and remember it so that we can use it when we try to connect
        // later.
        tunnelDatabasePort = address.getPort();
  
        LOGGER.info("Established tunneling session.  Port forwarding started on " + address.toInetSocketAddress());
        return session;
      } catch (final IOException e) {
        throw new RuntimeException(e);
      }
    }
  
    @Override
    public String toString() {
      return "SshTunnel{" +
          "hostKey=" + hostKey +
          ", portKey=" + portKey +
          ", tunnelMethod=" + tunnelMethod +
          ", tunnelHost='" + tunnelHost + '\'' +
          ", tunnelPort=" + tunnelPort +
          ", tunnelUser='" + tunnelUser + '\'' +
          ", remoteDatabaseHost='" + remoteDatabaseHost + '\'' +
          ", remoteDatabasePort=" + remoteDatabasePort +
          ", tunnelDatabasePort=" + tunnelDatabasePort +
          '}';
    }
  
  }
  