diff --git a/pkg/tools/cron.go b/pkg/tools/cron.go index 404dead70..a8ccdfdf5 100644 --- a/pkg/tools/cron.go +++ b/pkg/tools/cron.go @@ -154,11 +154,11 @@ func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult case "update": return t.updateJob(ctx, args) case "remove": - return t.removeJob(args) + return t.removeJob(ctx, args) case "enable": - return t.enableJob(args, true) + return t.enableJob(ctx, args, true) case "disable": - return t.enableJob(args, false) + return t.enableJob(ctx, args, false) default: return ErrorResult(fmt.Sprintf("unknown action: %s", action)) } @@ -377,12 +377,20 @@ func (t *CronTool) updateJob(ctx context.Context, args map[string]any) *ToolResu return SilentResult(fmt.Sprintf("Cron job updated:\n%s", formatCronJobJSON(updated))) } -func (t *CronTool) removeJob(args map[string]any) *ToolResult { +func (t *CronTool) removeJob(ctx context.Context, args map[string]any) *ToolResult { jobID, ok := args["job_id"].(string) if !ok || jobID == "" { return ErrorResult("job_id is required for remove") } + job, ok := t.cronService.GetJob(jobID) + if !ok { + return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) + } + if !t.canAccessJob(ctx, job) { + return ErrorResult(fmt.Sprintf("Job %s is not accessible from this channel", jobID)) + } + if t.cronService.RemoveJob(jobID) { return SilentResult(fmt.Sprintf("Cron job removed: %s", jobID)) } @@ -534,14 +542,18 @@ func (t *CronTool) canAccessJob(ctx context.Context, job *cron.CronJob) bool { if constants.IsInternalChannel(channel) { return true } + chatID := ToolChatID(ctx) if channel == "" || chatID == "" { return false } - if job.Payload.Command != "" { + if job.Payload.Channel != channel || job.Payload.To != chatID { return false } - return job.Payload.Channel == channel && job.Payload.To == chatID + if job.Payload.Command != "" { + return isCommandAllowedRemote(channel, chatID, t.commandAllowedRemotes) + } + return true } func formatCronJobJSON(job *cron.CronJob) string { @@ -552,14 +564,22 @@ func formatCronJobJSON(job *cron.CronJob) string { return string(data) } -func (t *CronTool) enableJob(args map[string]any, enable bool) *ToolResult { +func (t *CronTool) enableJob(ctx context.Context, args map[string]any, enable bool) *ToolResult { jobID, ok := args["job_id"].(string) if !ok || jobID == "" { return ErrorResult("job_id is required for enable/disable") } - job := t.cronService.EnableJob(jobID, enable) - if job == nil { + job, ok := t.cronService.GetJob(jobID) + if !ok { + return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) + } + if !t.canAccessJob(ctx, job) { + return ErrorResult(fmt.Sprintf("Job %s is not accessible from this channel", jobID)) + } + + updatedJob := t.cronService.EnableJob(jobID, enable) + if updatedJob == nil { return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) } @@ -567,7 +587,7 @@ func (t *CronTool) enableJob(args map[string]any, enable bool) *ToolResult { if !enable { status = "disabled" } - return SilentResult(fmt.Sprintf("Cron job '%s' %s", job.Name, status)) + return SilentResult(fmt.Sprintf("Cron job '%s' %s", updatedJob.Name, status)) } // ExecuteJob executes a cron job through the agent diff --git a/pkg/tools/cron_test.go b/pkg/tools/cron_test.go index 7118ca28a..7a9971d89 100644 --- a/pkg/tools/cron_test.go +++ b/pkg/tools/cron_test.go @@ -87,6 +87,28 @@ func parseCronJobResult(t *testing.T, result *ToolResult) cron.CronJob { return job } +func addTestCronJob(t *testing.T, tool *CronTool, name, channel, chatID, command string) *cron.CronJob { + t.Helper() + everyMS := int64(60_000) + job, err := tool.cronService.AddJob( + name, + cron.CronSchedule{Kind: "every", EveryMS: &everyMS}, + name+" message", + channel, + chatID, + ) + if err != nil { + t.Fatalf("AddJob() error: %v", err) + } + if command != "" { + job.Payload.Command = command + if err := tool.cronService.UpdateJob(job); err != nil { + t.Fatalf("UpdateJob() error: %v", err) + } + } + return job +} + // TestCronTool_CommandBlockedFromRemoteChannel verifies command scheduling is restricted by default. func TestCronTool_CommandBlockedFromRemoteChannel(t *testing.T) { tool := newTestCronTool(t) @@ -694,6 +716,296 @@ func TestCronTool_RemoteCannotAccessCommandJob(t *testing.T) { } } +func TestCronTool_AllowlistedRemoteCanAccessOwnCommandJob(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"telegram:chat-1"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + if listResult.IsError || !strings.Contains(listResult.ForLLM, job.ID) { + t.Fatalf("expected list to include own command job, got: %+v", listResult) + } + + getResult := tool.Execute(ctx, map[string]any{"action": "get", "job_id": job.ID}) + if getResult.IsError { + t.Fatalf("expected get to access own command job, got: %s", getResult.ForLLM) + } + got := parseCronJobResult(t, getResult) + if got.ID != job.ID || got.Payload.Command != "df -h" { + t.Fatalf("get returned wrong command job: %+v", got) + } + + updateResult := tool.Execute(ctx, map[string]any{ + "action": "update", + "job_id": job.ID, + "message": "updated command description", + }) + if updateResult.IsError { + t.Fatalf("expected update to access own command job, got: %s", updateResult.ForLLM) + } + updated, _ := tool.cronService.GetJob(job.ID) + if updated.Payload.Message != "updated command description" || updated.Payload.Command != "df -h" { + t.Fatalf("update returned wrong command payload: %+v", updated.Payload) + } +} + +func TestCronTool_AllowlistedRemoteCannotAccessOtherChatCommandJob(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"telegram"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-2", "df -h") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + if listResult.IsError || strings.Contains(listResult.ForLLM, job.ID) { + t.Fatalf("expected list to hide other chat command job, got: %+v", listResult) + } + + for _, action := range []string{"get", "update"} { + args := map[string]any{"action": action, "job_id": job.ID} + if action == "update" { + args["message"] = "changed" + } + result := tool.Execute(ctx, args) + if !result.IsError || !strings.Contains(result.ForLLM, "not accessible") { + t.Fatalf("expected %s to reject other chat command job, got: %+v", action, result) + } + } +} + +func TestCronTool_NonAllowlistedRemoteCannotAccessOwnCommandJob(t *testing.T) { + tool := newTestCronTool(t) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + if listResult.IsError || strings.Contains(listResult.ForLLM, job.ID) { + t.Fatalf("expected list to hide non-allowlisted command job, got: %+v", listResult) + } + + for _, action := range []string{"get", "update"} { + args := map[string]any{"action": action, "job_id": job.ID} + if action == "update" { + args["message"] = "changed" + } + result := tool.Execute(ctx, args) + if !result.IsError || !strings.Contains(result.ForLLM, "not accessible") { + t.Fatalf("expected %s to reject non-allowlisted command job, got: %+v", action, result) + } + } +} + +func TestCronTool_WildcardRemoteCanAccessOwnCommandJob(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"*"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + other := addTestCronJob(t, tool, "other", "telegram", "chat-2", "uptime") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + if listResult.IsError || !strings.Contains(listResult.ForLLM, job.ID) { + t.Fatalf("expected wildcard list to include own command job, got: %+v", listResult) + } + if strings.Contains(listResult.ForLLM, other.ID) { + t.Fatalf("wildcard list should still hide other chat job, got: %s", listResult.ForLLM) + } + + getResult := tool.Execute(ctx, map[string]any{"action": "get", "job_id": job.ID}) + if getResult.IsError { + t.Fatalf("expected wildcard get to access own command job, got: %s", getResult.ForLLM) + } +} + +func TestCronTool_InternalChannelCanAccessAllCommandJobs(t *testing.T) { + tool := newTestCronTool(t) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + ctx := WithToolContext(context.Background(), "cli", "direct") + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + if listResult.IsError || !strings.Contains(listResult.ForLLM, job.ID) { + t.Fatalf("expected internal list to include command job, got: %+v", listResult) + } + + getResult := tool.Execute(ctx, map[string]any{"action": "get", "job_id": job.ID}) + if getResult.IsError { + t.Fatalf("expected internal get to access command job, got: %s", getResult.ForLLM) + } + + updateResult := tool.Execute(ctx, map[string]any{ + "action": "update", + "job_id": job.ID, + "message": "internal update", + }) + if updateResult.IsError { + t.Fatalf("expected internal update to access command job, got: %s", updateResult.ForLLM) + } +} + +func TestCronTool_AllowlistedRemoteCanManageOwnCommandJob(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"telegram:chat-1"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + if action == "enable" { + tool.cronService.EnableJob(job.ID, false) + } + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if result.IsError { + t.Fatalf("expected %s to access own command job, got: %s", action, result.ForLLM) + } + + saved, ok := tool.cronService.GetJob(job.ID) + switch action { + case "remove": + if ok { + t.Fatalf("remove should delete own command job: %+v", saved) + } + case "enable": + if !ok || !saved.Enabled { + t.Fatalf("enable should enable own command job: %+v", saved) + } + case "disable": + if !ok || saved.Enabled { + t.Fatalf("disable should disable own command job: %+v", saved) + } + } + }) + } +} + +func TestCronTool_RemoteCannotManageOtherChatJob(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"telegram"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-2", "df -h") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if !result.IsError || !strings.Contains(result.ForLLM, "not accessible") { + t.Fatalf("expected %s to reject other chat job, got: %+v", action, result) + } + + saved, ok := tool.cronService.GetJob(job.ID) + if !ok { + t.Fatalf("%s should not remove other chat job", action) + } + if !saved.Enabled { + t.Fatalf("%s should not disable other chat job: %+v", action, saved) + } + }) + } +} + +func TestCronTool_RemoteCannotManageCommandJobUnlessAllowlisted(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + tool := newTestCronTool(t) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if !result.IsError || !strings.Contains(result.ForLLM, "not accessible") { + t.Fatalf("expected %s to reject non-allowlisted command job, got: %+v", action, result) + } + + saved, ok := tool.cronService.GetJob(job.ID) + if !ok { + t.Fatalf("%s should not remove non-allowlisted command job", action) + } + if !saved.Enabled { + t.Fatalf("%s should not disable non-allowlisted command job: %+v", action, saved) + } + }) + } +} + +func TestCronTool_InternalChannelCanManageAllJobs(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + tool := newTestCronTool(t) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + if action == "enable" { + tool.cronService.EnableJob(job.ID, false) + } + ctx := WithToolContext(context.Background(), "cli", "direct") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if result.IsError { + t.Fatalf("expected internal %s to access command job, got: %s", action, result.ForLLM) + } + + saved, ok := tool.cronService.GetJob(job.ID) + switch action { + case "remove": + if ok { + t.Fatalf("internal remove should delete command job: %+v", saved) + } + case "enable": + if !ok || !saved.Enabled { + t.Fatalf("internal enable should enable command job: %+v", saved) + } + case "disable": + if !ok || saved.Enabled { + t.Fatalf("internal disable should disable command job: %+v", saved) + } + } + }) + } +} + +func TestCronTool_RemoteCanManageOwnNonCommandJob(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + tool := newTestCronTool(t) + job := addTestCronJob(t, tool, "reminder", "telegram", "chat-1", "") + if action == "enable" { + tool.cronService.EnableJob(job.ID, false) + } + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if result.IsError { + t.Fatalf("expected %s to access own non-command job, got: %s", action, result.ForLLM) + } + }) + } +} + +func TestCronTool_WildcardRemoteCanManageOwnCommandJob(t *testing.T) { + for _, action := range []string{"remove", "enable", "disable"} { + t.Run(action, func(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Tools.Cron.CommandAllowedRemotes = []string{"*"} + tool := newTestCronToolWithConfig(t, cfg) + job := addTestCronJob(t, tool, "command", "telegram", "chat-1", "df -h") + if action == "enable" { + tool.cronService.EnableJob(job.ID, false) + } + other := addTestCronJob(t, tool, "other", "telegram", "chat-2", "uptime") + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + + result := tool.Execute(ctx, map[string]any{"action": action, "job_id": job.ID}) + if result.IsError { + t.Fatalf("expected wildcard %s to access own command job, got: %s", action, result.ForLLM) + } + + otherResult := tool.Execute(ctx, map[string]any{"action": action, "job_id": other.ID}) + if !otherResult.IsError || !strings.Contains(otherResult.ForLLM, "not accessible") { + t.Fatalf("wildcard %s should still reject other chat job, got: %+v", action, otherResult) + } + }) + } +} + func TestCronTool_CommandUpdateSafetyGates(t *testing.T) { t.Run("exec disabled", func(t *testing.T) { cfg := config.DefaultConfig()