runnable-sftp.go

package jobqueue

import (
	"fmt"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"io"
	"os"
)

// SFTPResult is a result of a sftp
type SFTPResult struct {
	FilesCopied []string
}

const (
	CredentialTypePassword = "password"
	CredentialTypeKey      = "key"
)

type Direction string

const (
	LocalToRemote Direction = "LocalToRemote"
	RemoteToLocal Direction = "RemoteToLocal"
)

type SFTPRunnable struct {
	Host              string
	Port              int
	User              string
	Insecure          bool
	Credential        string
	CredentialType    string
	HostKey           string
	SrcDir            string
	DstDir            string
	TransferDirection Direction
}

func (s *SFTPRunnable) Run() (RunResult[SFTPResult], error) {

	var authMethod ssh.AuthMethod

	// Auth
	switch s.CredentialType {
	case CredentialTypePassword:
		authMethod = ssh.Password(s.Credential)
	case CredentialTypeKey:
		key, err := ssh.ParsePrivateKey([]byte(s.Credential))
		if err != nil {
			return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
		}
		authMethod = ssh.PublicKeys(key)
	default:
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, ErrUnsupportedCredentialType
	}

	var hkCallback ssh.HostKeyCallback

	if s.HostKey != "" {

		hostkeyBytes := []byte(s.HostKey)
		hostKey, err := ssh.ParsePublicKey(hostkeyBytes)
		if err != nil {
			return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
		}

		hkCallback = ssh.FixedHostKey(hostKey)
	} else {
		if s.Insecure {
			hkCallback = ssh.InsecureIgnoreHostKey()
		} else {
			hkCallback = ssh.FixedHostKey(nil)
		}
	}

	config := &ssh.ClientConfig{
		User: s.User,
		Auth: []ssh.AuthMethod{
			authMethod,
		},
		HostKeyCallback: hkCallback,
	}

	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", s.Host, s.Port), config)
	if err != nil {
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
	}
	defer client.Close()

	sftpClient, err := sftp.NewClient(client)
	if err != nil {
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
	}
	defer sftpClient.Close()

	var filesCopied []string

	switch s.TransferDirection {
	case LocalToRemote:
		filesCopied, err = s.copyLocalToRemote(sftpClient)
	case RemoteToLocal:
		filesCopied, err = s.copyRemoteToLocal(sftpClient)
	default:
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, ErrUnsupportedTransferDirection
	}

	if err != nil {
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
	}

	if err != nil {
		return RunResult[SFTPResult]{Status: ResultStatusFailed}, err
	}

	return RunResult[SFTPResult]{Status: ResultStatusSuccess, Data: SFTPResult{FilesCopied: filesCopied}}, nil
}

func copyFile(src io.Reader, dst io.Writer) error {
	_, err := io.Copy(dst, src)
	return err
}

func (s *SFTPRunnable) copyLocalToRemote(sftpClient *sftp.Client) ([]string, error) {

	var filesCopied []string

	// create destination directory
	err := sftpClient.MkdirAll(s.DstDir)
	if err != nil {
		return nil, err
	}

	// copy files
	files, err := os.ReadDir(s.SrcDir)
	if err != nil {
		return nil, err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		srcFile, err := os.Open(fmt.Sprintf("%s/%s", s.SrcDir, file.Name()))
		if err != nil {
			return nil, err
		}
		dstFile, err := sftpClient.Create(fmt.Sprintf("%s/%s", s.DstDir, file.Name()))
		if err != nil {
			_ = srcFile.Close()
			return nil, err
		}
		err = copyFile(srcFile, dstFile)
		_ = srcFile.Close()
		_ = dstFile.Close()
		if err != nil {
			return nil, err
		}

		filesCopied = append(filesCopied, fmt.Sprintf("%s/%s", s.DstDir, file.Name()))
	}

	return filesCopied, nil
}

func (s *SFTPRunnable) copyRemoteToLocal(sftpClient *sftp.Client) ([]string, error) {

	var filesCopied []string

	// create destination directory
	err := os.MkdirAll(s.DstDir, 0755)
	if err != nil {
		return nil, err
	}

	// copy files
	files, err := sftpClient.ReadDir(s.SrcDir)
	if err != nil {
		return nil, err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		srcFile, err := sftpClient.Open(fmt.Sprintf("%s/%s", s.SrcDir, file.Name()))
		if err != nil {
			return nil, err
		}
		dstFile, err := os.Create(fmt.Sprintf("%s/%s", s.DstDir, file.Name()))
		if err != nil {
			_ = srcFile.Close()
			return nil, err
		}
		err = copyFile(srcFile, dstFile)
		_ = srcFile.Close()
		_ = dstFile.Close()
		if err != nil {
			return nil, err
		}

		filesCopied = append(filesCopied, fmt.Sprintf("%s/%s", s.DstDir, file.Name()))
	}

	return filesCopied, nil
}