wraith/internal/sftp/service.go

package sftp

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	"github.com/pkg/sftp"
)

const MaxEditFileSize = 5 * 1024 * 1024 // 5MB

// FileEntry represents a file or directory in a remote filesystem.
type FileEntry struct {
	Name        string `json:"name"`
	Path        string `json:"path"`
	Size        int64  `json:"size"`
	IsDir       bool   `json:"isDir"`
	Permissions string `json:"permissions"`
	ModTime     string `json:"modTime"`
}

// SFTPService manages SFTP clients keyed by session ID.
type SFTPService struct {
	clients map[string]*sftp.Client
	mu      sync.RWMutex
}

// NewSFTPService creates a new SFTPService.
func NewSFTPService() *SFTPService {
	return &SFTPService{
		clients: make(map[string]*sftp.Client),
	}
}

// RegisterClient stores an SFTP client for a session.
func (s *SFTPService) RegisterClient(sessionID string, client *sftp.Client) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.clients[sessionID] = client
}

// RemoveClient removes and closes an SFTP client.
func (s *SFTPService) RemoveClient(sessionID string) {
	s.mu.Lock()
	client, ok := s.clients[sessionID]
	if ok {
		delete(s.clients, sessionID)
	}
	s.mu.Unlock()

	if ok && client != nil {
		client.Close()
	}
}

// getClient returns the SFTP client for a session or an error if not found.
func (s *SFTPService) getClient(sessionID string) (*sftp.Client, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	client, ok := s.clients[sessionID]
	if !ok {
		return nil, fmt.Errorf("no SFTP client for session %s", sessionID)
	}
	return client, nil
}

// SortEntries sorts file entries with directories first, then alphabetically by name.
func SortEntries(entries []FileEntry) {
	sort.Slice(entries, func(i, j int) bool {
		if entries[i].IsDir != entries[j].IsDir {
			return entries[i].IsDir
		}
		return strings.ToLower(entries[i].Name) < strings.ToLower(entries[j].Name)
	})
}

// fileInfoToEntry converts an os.FileInfo and its path into a FileEntry.
func fileInfoToEntry(info os.FileInfo, path string) FileEntry {
	return FileEntry{
		Name:        info.Name(),
		Path:        path,
		Size:        info.Size(),
		IsDir:       info.IsDir(),
		Permissions: info.Mode().Perm().String(),
		ModTime:     info.ModTime().UTC().Format("2006-01-02T15:04:05Z"),
	}
}

// List returns directory contents sorted (dirs first, then files alphabetically).
func (s *SFTPService) List(sessionID string, path string) ([]FileEntry, error) {
	client, err := s.getClient(sessionID)
	if err != nil {
		return nil, err
	}

	infos, err := client.ReadDir(path)
	if err != nil {
		return nil, fmt.Errorf("read directory %s: %w", path, err)
	}

	entries := make([]FileEntry, 0, len(infos))
	for _, info := range infos {
		entryPath := path
		if !strings.HasSuffix(entryPath, "/") {
			entryPath += "/"
		}
		entryPath += info.Name()
		entries = append(entries, fileInfoToEntry(info, entryPath))
	}

	SortEntries(entries)
	return entries, nil
}

// ReadFile reads a file (max 5MB). Returns content as string.
func (s *SFTPService) ReadFile(sessionID string, path string) (string, error) {
	client, err := s.getClient(sessionID)
	if err != nil {
		return "", err
	}

	info, err := client.Stat(path)
	if err != nil {
		return "", fmt.Errorf("stat %s: %w", path, err)
	}

	if info.IsDir() {
		return "", fmt.Errorf("%s is a directory", path)
	}

	if info.Size() > MaxEditFileSize {
		return "", fmt.Errorf("file %s is %d bytes, exceeds max edit size of %d bytes", path, info.Size(), MaxEditFileSize)
	}

	f, err := client.Open(path)
	if err != nil {
		return "", fmt.Errorf("open %s: %w", path, err)
	}
	defer f.Close()

	data, err := io.ReadAll(f)
	if err != nil {
		return "", fmt.Errorf("read %s: %w", path, err)
	}

	return string(data), nil
}

// WriteFile writes content to a file.
func (s *SFTPService) WriteFile(sessionID string, path string, content string) error {
	client, err := s.getClient(sessionID)
	if err != nil {
		return err
	}

	f, err := client.Create(path)
	if err != nil {
		return fmt.Errorf("create %s: %w", path, err)
	}
	defer f.Close()

	if _, err := f.Write([]byte(content)); err != nil {
		return fmt.Errorf("write %s: %w", path, err)
	}

	return nil
}

// Mkdir creates a directory.
func (s *SFTPService) Mkdir(sessionID string, path string) error {
	client, err := s.getClient(sessionID)
	if err != nil {
		return err
	}

	if err := client.Mkdir(path); err != nil {
		return fmt.Errorf("mkdir %s: %w", path, err)
	}
	return nil
}

// Delete removes a file or empty directory.
func (s *SFTPService) Delete(sessionID string, path string) error {
	client, err := s.getClient(sessionID)
	if err != nil {
		return err
	}

	info, err := client.Stat(path)
	if err != nil {
		return fmt.Errorf("stat %s: %w", path, err)
	}

	if info.IsDir() {
		if err := client.RemoveDirectory(path); err != nil {
			return fmt.Errorf("remove directory %s: %w", path, err)
		}
	} else {
		if err := client.Remove(path); err != nil {
			return fmt.Errorf("remove %s: %w", path, err)
		}
	}

	return nil
}

// Rename renames/moves a file.
func (s *SFTPService) Rename(sessionID string, oldPath, newPath string) error {
	client, err := s.getClient(sessionID)
	if err != nil {
		return err
	}

	if err := client.Rename(oldPath, newPath); err != nil {
		return fmt.Errorf("rename %s to %s: %w", oldPath, newPath, err)
	}
	return nil
}

// Stat returns info about a file/directory.
func (s *SFTPService) Stat(sessionID string, path string) (*FileEntry, error) {
	client, err := s.getClient(sessionID)
	if err != nil {
		return nil, err
	}

	info, err := client.Stat(path)
	if err != nil {
		return nil, fmt.Errorf("stat %s: %w", path, err)
	}

	entry := fileInfoToEntry(info, path)
	return &entry, nil
}