goexplore/internal/protocols/sftp/sftp.go

package sftp

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"goexplore/internal/config"
	"goexplore/internal/explorer"
)

type SFTPExplorer struct {
	cfg    *config.ConnectionConfig
	secret string
	client *sftp.Client
	conn   *ssh.Client
}

func New(c *config.ConnectionConfig, secret string) *SFTPExplorer {
	return &SFTPExplorer{cfg: c, secret: secret}
}

func (e *SFTPExplorer) Connect() error {
	port := e.cfg.Port
	if port == 0 {
		port = 22
	}
	
	var authMethods []ssh.AuthMethod
	
	secretStr := strings.TrimSpace(e.secret)
	// Handle case where newlines might be escaped by mistake
	secretStr = strings.ReplaceAll(secretStr, "\\n", "\n")

	if strings.Contains(secretStr, "-----BEGIN") {
		signer, err := ssh.ParsePrivateKey([]byte(secretStr))
		if err != nil {
			return fmt.Errorf("failed to parse private key: %w", err)
		}
		authMethods = append(authMethods, ssh.PublicKeys(signer))
	} else {
		authMethods = append(authMethods, ssh.Password(e.secret))
	}

	// Always add KeyboardInteractive to support Duo MFA and servers that require interactive auth.
	authMethods = append(authMethods, ssh.KeyboardInteractive(func(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
		answers = make([]string, len(questions))
		for i, q := range questions {
			qLower := strings.ToLower(q)
			if strings.Contains(secretStr, "-----BEGIN") {
				// If using a private key, we don't have a password.
				// If it's a Duo prompt asking for an option, "1" is typically Duo Push.
				if strings.Contains(qLower, "duo push") || strings.Contains(qLower, "passcode or option") {
					answers[i] = "1"
				} else {
					answers[i] = ""
				}
			} else {
				// Otherwise, respond with the password to keyboard interactive prompts
				answers[i] = e.secret
			}
		}
		return answers, nil
	}))

	config := &ssh.ClientConfig{
		User: e.cfg.Username,
		Auth: authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", e.cfg.Host, port), config)
	if err != nil {
		return err
	}
	e.conn = conn

	client, err := sftp.NewClient(conn)
	if err != nil {
		conn.Close()
		return err
	}
	e.client = client
	return nil
}

func (e *SFTPExplorer) Disconnect() error {
	if e.client != nil {
		e.client.Close()
	}
	if e.conn != nil {
		e.conn.Close()
	}
	return nil
}

func (e *SFTPExplorer) ListDir(path string) ([]explorer.FileEntry, error) {
	if path == "" {
		path = "."
	}
	files, err := e.client.ReadDir(path)
	if err != nil {
		return nil, err
	}

	var entries []explorer.FileEntry
	for _, f := range files {
		entries = append(entries, explorer.FileEntry{
			Name:        f.Name(),
			Path:        filepath.ToSlash(filepath.Join(path, f.Name())),
			Size:        f.Size(),
			Modified:    f.ModTime().Format(time.RFC3339),
			IsDir:       f.IsDir(),
			Permissions: f.Mode().String(),
		})
	}
	return entries, nil
}

func (e *SFTPExplorer) Stat(path string) (explorer.FileEntry, error) {
	if path == "" {
		path = "."
	}
	f, err := e.client.Stat(path)
	if err != nil {
		return explorer.FileEntry{}, err
	}
	return explorer.FileEntry{
		Name:        f.Name(),
		Path:        filepath.ToSlash(path),
		Size:        f.Size(),
		Modified:    f.ModTime().Format(time.RFC3339),
		IsDir:       f.IsDir(),
		Permissions: f.Mode().String(),
	}, nil
}

func (e *SFTPExplorer) MkDir(path string) error {
	return e.client.MkdirAll(path)
}

func (e *SFTPExplorer) removeAll(path string) error {
	fi, err := e.client.Stat(path)
	if err != nil {
		return err
	}
	if !fi.IsDir() {
		return e.client.Remove(path)
	}

	files, err := e.client.ReadDir(path)
	if err != nil {
		return err
	}

	for _, f := range files {
		if f.Name() == "." || f.Name() == ".." {
			continue
		}
		if err := e.removeAll(filepath.Join(path, f.Name())); err != nil {
			return err
		}
	}

	return e.client.RemoveDirectory(path)
}

func (e *SFTPExplorer) Delete(path string) error {
	return e.removeAll(path)
}

func (e *SFTPExplorer) Rename(src, dst string) error {
	return e.client.Rename(src, dst)
}

func (e *SFTPExplorer) ReadFile(path string) (io.ReadCloser, error) {
	return e.client.Open(path)
}

func (e *SFTPExplorer) WriteFile(path string, r io.Reader, size int64) error {
	f, err := e.client.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, r)
	return err
}

func (e *SFTPExplorer) Checksum(path string) (string, error) {
	r, err := e.ReadFile(path)
	if err != nil {
		return "", err
	}
	defer r.Close()
	h := md5.New()
	if _, err := io.Copy(h, r); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}