exec/ssh/rootfs.go

package ssh

import (
	"archive/tar"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"sirherobrine23.com.br/go-bds/exec"
)

var _ exec.Rootfs = (*Rootfs)(nil)

func sshRegister(args ...any) (*Rootfs, error) {
	var sshClient *ssh.Client
	var err error

	if len(args) > 0 {
		switch v := args[0].(type) {
		case *ssh.Client:
			sshClient = v
		case string:
			sshConnection := v
			if len(args) > 1 {
				switch v := args[1].(type) {
				case *ssh.ClientConfig:
					if sshClient, err = ssh.Dial("tcp", sshConnection, v); err != nil {
						return nil, err
					}
				case string:
					if len(args) <= 2 {
						return nil, fmt.Errorf("require ssh password")
					}

					sshUser := v
					sshPassword, ok := args[2].(string)
					if !ok {
						return nil, fmt.Errorf("require ssh password string")
					}

					sshClient, err = ssh.Dial("tcp", sshConnection, &ssh.ClientConfig{
						User:            sshUser,
						HostKeyCallback: ssh.InsecureIgnoreHostKey(),
						Auth:            []ssh.AuthMethod{ssh.Password(sshPassword)},
					})

					if err != nil {
						return nil, err
					}
				default:
					return nil, fmt.Errorf("require ssh user")
				}
			}
		default:
			return nil, fmt.Errorf("require ssh client config or ssh connection string")
		}
	}

	return &Rootfs{Client: sshClient}, nil
}

var _ = exec.RegisterRootfs("ssh", sshRegister)

func NewClient(dial string, config *ssh.ClientConfig) (*Rootfs, error) {
	c, err := ssh.Dial("tcp", dial, config)
	if err != nil {
		return nil, err
	}
	return &Rootfs{Client: c}, nil
}

func NewClientPassword(dial string, user, pass string) (*Rootfs, error) {
	return sshRegister(dial, user, pass)
}

type Rootfs struct {
	Path string

	Client *ssh.Client
}

func (rootfs *Rootfs) Close() error {
	return rootfs.Client.Close()
}

func (rootfs *Rootfs) Tar(w io.Writer) error {
	sftpClient, err := sftp.NewClient(rootfs.Client)
	if err != nil {
		return err
	}
	defer sftpClient.Close()
	t := tar.NewWriter(w)
	defer t.Close()

	walk := sftpClient.Walk(rootfs.Path)
	for walk.Step() {
		if err = walk.Err(); err != nil {
			return err
		}
		name, info := walk.Path(), walk.Stat()
		realPath, _ := sftpClient.RealPath(name)

		h, err := tar.FileInfoHeader(info, realPath)
		if err != nil {
			return err
		}
		h.Name = name

		// Write header
		if err = t.WriteHeader(h); err != nil {
			return err
		}

		if info.IsDir() || !info.Mode().Type().IsRegular() {
			continue
		}

		f, err := sftpClient.Open(name)
		if err != nil {
			return err
		}
		defer f.Close()

		if _, err = io.CopyN(t, f, info.Size()); err != nil {
			return err
		}
	}

	return nil
}

func (rootfs *Rootfs) FromTar(r io.Reader) error {
	sftpClient, err := sftp.NewClient(rootfs.Client)
	if err != nil {
		return err
	}
	defer sftpClient.Close()

	t := tar.NewReader(r)
	for {
		h, err := t.Next()
		switch err {
		case nil:
		case io.EOF:
			return nil
		default:
			return err
		}
		rootName := filepath.Join(rootfs.Path, h.Name)

		switch h.FileInfo().Mode().Type() {
		case fs.ModeNamedPipe:
		case fs.ModeSocket:
		case fs.ModeDevice:
		case fs.ModeCharDevice:
		case fs.ModeAppend:
		case fs.ModeExclusive:
		case fs.ModeTemporary:
		case fs.ModeSetuid:
		case fs.ModeSetgid:
		case fs.ModeSticky:
		case fs.ModeIrregular:
		case fs.ModeDir:
			if err := sftpClient.Mkdir(rootName); err != nil {
				return err
			} else if err = sftpClient.Chmod(rootName, h.FileInfo().Mode().Perm()); err != nil {
				return err
			}
		case fs.ModeSymlink:
			if err := sftpClient.Symlink(h.Linkname, rootName); err != nil {
				return err
			}
		default:
			f, err := sftpClient.OpenFile(rootName, os.O_CREATE|os.O_WRONLY|os.O_TRUNC)
			if err != nil {
				return err
			}

			// Copy Content body
			_, err = io.Copy(f, t)
			f.Chmod(h.FileInfo().Mode().Perm())
			f.Chown(h.Uid, h.Gid)
			f.Close()
			if err != nil {
				return err
			}
		}
	}
}

func (rootfs *Rootfs) NewProcess() (exec.Proc, error) {
	session, err := rootfs.Client.NewSession()
	if err != nil {
		return nil, err
	}
	return &Session{
		Session: session,
		Client:  rootfs.Client,
	}, nil
}