public static ClientSession getConnectSession()

in hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelper.java [109:205]
88 lines of code
23 McCabe index (conditional complexity)

    public static ClientSession getConnectSession(SshProtocol sshProtocol, int timeout, boolean reuseConnection, boolean useProxy)
            throws IOException, GeneralSecurityException {
        CacheIdentifier identifier = CacheIdentifier.builder()
                                                    .ip(sshProtocol.getHost()).port(sshProtocol.getPort())
                                                    .username(sshProtocol.getUsername()).password(sshProtocol.getPassword())
                                                    .build();
        ClientSession clientSession = null;
        // When using ProxyJump, force connection reuse:
        // Apache MINA SSHD will pass the proxy password error to the target host in proxy scenarios, causing the first connection to fail.
        // Reusing connections can skip duplicate authentication and avoid this problem.
        if (reuseConnection || useProxy) {
            Optional<AbstractConnection<?>> cacheOption = CONNECTION_COMMON_CACHE.getCache(identifier, true);
            if (cacheOption.isPresent()) {
                SshConnect sshConnect = (SshConnect) cacheOption.get();
                clientSession = sshConnect.getConnection();
                try {
                    if (clientSession == null || clientSession.isClosed() || clientSession.isClosing()) {
                        clientSession = null;
                        CONNECTION_COMMON_CACHE.removeCache(identifier);
                    }
                } catch (Exception e) {
                    log.warn(e.getMessage());
                    clientSession = null;
                    CONNECTION_COMMON_CACHE.removeCache(identifier);
                }
            }
            if (clientSession != null) {
                return clientSession;
            }
        }
        SshClient sshClient = CommonSshClient.getSshClient();
        HostConfigEntry proxyConfig = new HostConfigEntry();
        if (useProxy && StringUtils.hasText(sshProtocol.getProxyHost())) {
            String proxySpec = String.format("%s@%s:%d", sshProtocol.getProxyUsername(), sshProtocol.getProxyHost(), Integer.parseInt(sshProtocol.getProxyPort()));
            proxyConfig.setHostName(sshProtocol.getHost());
            proxyConfig.setHost(sshProtocol.getHost());
            proxyConfig.setPort(Integer.parseInt(sshProtocol.getPort()));
            proxyConfig.setUsername(sshProtocol.getUsername());
            proxyConfig.setProxyJump(proxySpec);

            // Apache SSHD requires the password for the proxy to be preloaded into the sshClient instance before connecting
            if (StringUtils.hasText(sshProtocol.getProxyPassword())) {
                sshClient.addPasswordIdentity(sshProtocol.getProxyPassword());
                log.debug("Loaded proxy server password authentication: {}@{}", sshProtocol.getProxyUsername(), sshProtocol.getProxyHost());
            }
            if (StringUtils.hasText(sshProtocol.getProxyPrivateKey())) {
                proxyConfig.setIdentities(List.of(sshProtocol.getProxyPrivateKey()));
                log.debug("Proxy private key loaded into HostConfigEntry");
            }
        }

        if (useProxy && StringUtils.hasText(sshProtocol.getProxyHost())) {
            try {
                clientSession = sshClient.connect(proxyConfig)
                                         .verify(timeout, TimeUnit.MILLISECONDS).getSession();
            }
            finally {
                sshClient.removePasswordIdentity(sshProtocol.getProxyPassword());
            }
        } else {
            clientSession = sshClient.connect(sshProtocol.getUsername(), sshProtocol.getHost(), Integer.parseInt(sshProtocol.getPort()))
                                     .verify(timeout, TimeUnit.MILLISECONDS).getSession();
        }

        if (StringUtils.hasText(sshProtocol.getPassword())) {
            clientSession.addPasswordIdentity(sshProtocol.getPassword());
        } else if (StringUtils.hasText(sshProtocol.getPrivateKey())) {
            var resourceKey = PrivateKeyUtils.writePrivateKey(sshProtocol.getHost(), sshProtocol.getPrivateKey());
            try (InputStream keyStream = new FileInputStream(resourceKey)) {
                FilePasswordProvider passwordProvider = (session, resource, index) -> {
                    if (StringUtils.hasText(sshProtocol.getPrivateKeyPassphrase())) {
                        return sshProtocol.getPrivateKeyPassphrase();
                    }
                    return null;
                };
                Iterable<KeyPair> keyPairs = SecurityUtils.loadKeyPairIdentities(null, () -> resourceKey, keyStream, passwordProvider);
                if (keyPairs != null) {
                    keyPairs.forEach(clientSession::addPublicKeyIdentity);
                } else {
                    log.error("Failed to load private key pairs from: {}", resourceKey);
                }
            } catch (IOException e) {
                log.error("Error reading private key file: {}", e.getMessage());
            }
        }  // else auth with localhost private public key certificates

        // auth
        if (!clientSession.auth().verify(timeout, TimeUnit.MILLISECONDS).isSuccess()) {
            clientSession.close();
            throw new IllegalArgumentException("ssh auth failed.");
        }
        if (reuseConnection || useProxy) {
            SshConnect sshConnect = new SshConnect(clientSession);
            CONNECTION_COMMON_CACHE.addCache(identifier, sshConnect);
        }
        return clientSession;
    }