All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.wl4g.infra.common.cli.ssh.SshjHelper Maven / Gradle / Ivy

There is a newer version: 3.1.72
Show newest version
/*
 * Copyright 2017 ~ 2025 the original author or authors. James Wong 
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.wl4g.infra.common.cli.ssh;

import static com.wl4g.infra.common.io.ByteStreamUtils.readFullyToString;
import static com.wl4g.infra.common.lang.Assert2.hasText;
import static com.wl4g.infra.common.lang.Assert2.hasTextOf;
import static com.wl4g.infra.common.lang.Assert2.isTrueOf;
import static com.wl4g.infra.common.lang.Assert2.notNull;
import static com.wl4g.infra.common.lang.Assert2.notNullOf;
import static java.util.Objects.isNull;
import static java.util.Objects.nonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

import java.io.File;
import java.io.IOException;

import javax.validation.constraints.NotNull;

import com.google.common.annotations.Beta;
import com.wl4g.infra.common.cli.ssh.SshjHelper.CommandSessionWrapper;
import com.wl4g.infra.common.function.CallbackFunction;
import com.wl4g.infra.common.function.ProcessFunction;

import lombok.CustomLog;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.connection.channel.direct.Session;
import net.schmizz.sshj.connection.channel.direct.Session.Command;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import net.schmizz.sshj.userauth.keyprovider.KeyProvider;
import net.schmizz.sshj.xfer.FileSystemFile;
import net.schmizz.sshj.xfer.scp.SCPFileTransfer;
import net.schmizz.sshj.xfer.scp.ScpCommandLine;

/**
 * SSHJ based SSH2 tools.
 *
 * @author James Wong James Wong 
 * @version v1.0 2019年5月24日
 * @since
 */
@Beta
@CustomLog
public class SshjHelper extends SshHelperBase {

    private static final SshjHelper DEFAULT = new SshjHelper();

    public static SshjHelper getInstance() {
        return DEFAULT;
    }

    // --- Transfer files. ---

    @Override
    public void scpGetFile(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            File localFile,
            String remoteFilePath) throws Exception {
        notNull(localFile, "Transfer localFile must not be null.");
        hasText(remoteFilePath, "Transfer remoteDir can't empty.");

        log.debug("SSH2 transfer file from {} to {}@{}:{}", localFile.getAbsolutePath(), user, host, remoteFilePath);
        try {
            // Transfer get file.
            doScpTransfer(host, port, user, pemPrivateKey, password, scp -> {
                scp.download(remoteFilePath, new FileSystemFile(localFile));
            });

            log.debug("SCP get transfered: '{}' from '{}@{}:{}'", localFile.getAbsolutePath(), user, host, remoteFilePath);
        } catch (IOException e) {
            throw e;
        }
    }

    @Override
    public void scpPutFile(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            File localFile,
            String remoteDir) throws Exception {
        notNullOf(localFile, "localFile");
        hasTextOf(remoteDir, "remoteDir");

        log.debug("SSH2 transfer file from {} to {}@{}:{}", localFile.getAbsolutePath(), user, host, remoteDir);
        try {
            // Transfer send file.
            doScpTransfer(host, port, user, pemPrivateKey, password, scp -> {
                // scp.upload(new FileSystemFile(localFile), remoteDir);
                scp.newSCPUploadClient().copy(new FileSystemFile(localFile), remoteDir, ScpCommandLine.EscapeMode.NoEscape);
            });

            log.debug("SCP put transfered: '{}' to '{}@{}:{}'", localFile.getAbsolutePath(), user, host, remoteDir);
        } catch (IOException e) {
            throw e;
        }
    }

    /**
     * Perform file transfer with remote host, including scp.put/upload or
     * scp.get/download.
     * 
     * @param host
     * @param user
     * @param pemPrivateKey
     * @param processor
     * @throws IOException
     */
    @Override
    protected void doScpTransfer(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            CallbackFunction processor) throws Exception {
        hasTextOf(host, "host");
        isTrueOf(port >= 1, "port>=1");
        hasTextOf(user, "user");
        notNullOf(processor, "processor");

        // Fallback uses the local current user private key by default.
        if (isNull(pemPrivateKey)) {
            pemPrivateKey = getDefaultLocalUserPrivateKey();
        }
        notNull(pemPrivateKey, "Transfer pemPrivateKey can't null.");

        SSHClient ssh = null;
        try {
            ssh = new SSHClient();
            ssh.addHostKeyVerifier(new PromiscuousVerifier());
            ssh.connect(host, port);
            KeyProvider keyProvider = ssh.loadKeys(new String(pemPrivateKey), null, null);
            ssh.authPublickey(user, keyProvider);

            SCPFileTransfer scpFileTransfer = ssh.newSCPFileTransfer();

            // Transfer file(put/get).
            processor.process(scpFileTransfer);
        } catch (Exception e) {
            throw e;
        } finally {
            try {
                if (nonNull(ssh)) {
                    ssh.disconnect();
                    ssh.close();
                }
            } catch (Exception e) {
                log.error("Failed to closing ssh client.", e);
            }
        }
    }

    // --- Execution commands. ---

    public SSHExecResult execWaitForResponse(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            String command,
            long timeoutMs) throws Exception {
        return execWaitForComplete(host, port, user, pemPrivateKey, password, command, s -> {
            Session.Command cmd = s.getCommand();
            String message = null, errmsg = null;
            if (nonNull(cmd.getInputStream())) {
                message = readFullyToString(cmd.getInputStream());
            }
            if (nonNull(cmd.getErrorStream())) {
                errmsg = readFullyToString(cmd.getErrorStream());
            }
            return new SSHExecResult(nonNull(cmd.getExitSignal()) ? cmd.getExitSignal().toString() : null, cmd.getExitStatus(),
                    message, errmsg);
        }, timeoutMs);
    }

    @Override
    public  T execWaitForComplete(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            String command,
            @NotNull ProcessFunction processor,
            long timeoutMs) throws Exception {
        return doExecCommand(host, port, user, pemPrivateKey, password, command, s -> {
            // Wait for completed by condition.
            s.getCommand().join(timeoutMs, MILLISECONDS);
            return processor.process(s);
        });
    }

    @Override
    public  T doExecCommand(
            String host,
            int port,
            String user,
            char[] pemPrivateKey,
            String password,
            String command,
            ProcessFunction processor) throws Exception {
        hasTextOf(host, "host");
        isTrueOf(port >= 1, "port>=1");
        hasTextOf(user, "user");
        notNullOf(processor, "processor");

        // Fallback uses the local current user private key by default.
        if (isNull(pemPrivateKey) && isNull(password)) {
            pemPrivateKey = getDefaultLocalUserPrivateKey();
        }

        SSHClient ssh = null;
        Session session = null;
        Session.Command cmd = null;
        try {
            ssh = new SSHClient();
            ssh.addHostKeyVerifier(new PromiscuousVerifier());
            ssh.connect(host);

            if (nonNull(pemPrivateKey)) {
                KeyProvider keyProvider = ssh.loadKeys(new String(pemPrivateKey), null, null);
                ssh.authPublickey(user, keyProvider);
            } else {
                notNullOf(password, "password");
                ssh.authPassword(user, password);
            }
            session = ssh.startSession();

            // Note: temporarily load according to the priority of user
            // environment > global environment, ignoring errors caused by linux
            // OS differences (such as ubuntu without /etc/bashrc file).
            command = ". /etc/profile; . /etc/bashrc; . /etc/bash.bashrc; . ~/.profile; . ~/.bashrc; " + command;
            cmd = session.exec(command);

            return processor.process(new CommandSessionWrapper(session, cmd));
        } catch (Exception e) {
            throw e;
        } finally {
            try {
                if (nonNull(session)) {
                    session.close();
                }
            } catch (Exception e) {
                log.error("Closing sshj session failure", e);
            }
            try {
                if (nonNull(ssh)) {
                    ssh.disconnect();
                    ssh.close();
                }
            } catch (Exception e) {
                log.error("Closing sshj client failure", e);
            }
        }
    }

    // --- Tool function's. ---

    @Override
    public SSHKeyPair generateKeypair(AlgorithmType type, String comment) throws Exception {
        throw new UnsupportedOperationException();
    }

    /**
     * {@link CommandSessionWrapper}
     *
     * @author James Wong 
     * @version v1.0 2020-10-09
     * @since
     */
    public static class CommandSessionWrapper {
        private final Session session;
        private final Session.Command command;

        public CommandSessionWrapper(Session session, Command command) {
            super();
            this.session = session;
            this.command = command;
        }

        public Session getSession() {
            return session;
        }

        public Session.Command getCommand() {
            return command;
        }

    }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy