diff --git a/README.md b/README.md index c590a51..bf2ab16 100644 --- a/README.md +++ b/README.md @@ -30,8 +30,11 @@ bwh node add production --api-key --veid # Basic operations bwh info # View server details +bwh rate-limit # Check API rate limit status bwh start/stop/restart # Power management bwh usage --period 7d # Check usage statistics +bwh abuse suspensions # Show suspension details +bwh notifications list # Show notification preferences bwh snapshot create "backup-name" # Create snapshots bwh iso images # List available ISO images bwh iso mount ubuntu-20.04.iso # Mount ISO for rescue/install @@ -98,12 +101,16 @@ backups, err := c.ListBackups(ctx) **Server Management**: `GetServiceInfo`, `GetLiveServiceInfo`, `Start`, `Stop`, `Restart`, `Kill`, `SetHostname`, `ReinstallOS`, `ResetRootPassword`, `MountISO`, `UnmountISO` -**Monitoring**: `GetRawUsageStats`, `GetBasicServiceInfo`, audit log access +**Monitoring**: `GetRawUsageStats`, `GetAuditLog`, `GetRateLimitStatus` **Backup & Recovery**: `CreateSnapshot`, `RestoreSnapshot`, `DeleteSnapshot`, backup management **Migration**: `GetMigrateLocations`, `StartMigration` (use `StartMigrationWithTimeout` for custom timeouts) +**Security & Abuse**: `GetSuspensionDetails`, `GetPolicyViolations` + +**Notifications**: `GetNotificationPreferences` + **Network**: SSH key management, IP/reverse DNS configuration, IPv6 subnet management, private IPv4 management *Complete API reference*: View [pkg/client documentation](./pkg/client) or run `go doc github.com/strahe/bwh/pkg/client` for all available methods. @@ -187,7 +194,7 @@ Add to your Continue configuration: ### Configuration Notes - **Custom Config**: Use `--config /path/to/config.yaml` to specify a config file -- **Multiple Instances**: The server automatically uses your default instance from configuration +- **Multiple Instances**: The server uses the configured default instance, or `--instance ` as the MCP session default. Tool-level `instance` arguments override it. - **Integration**: Add to existing MCP config files without replacing other servers ### Available MCP Tools (Read-only) @@ -199,6 +206,14 @@ Add to your Continue configuration: - **backup_list**: List backups (`instance?`, `os_contains?`, `since?`, `until?`, `sort_by?`, `order?`, `limit?`) - **vps_audit_get**: Get audit logs (`instance?`, `since?`, `until?`, `limit?`, `ip_contains?`, `type?`) - **iso_list**: List available and mounted ISO images (`instance?`) +- **ssh_keys_get**: Get SSH public keys (`instance?`, `full?`) +- **os_templates_get**: List OS templates (`instance?`) +- **rate_limit_get**: Get API rate limit status (`instance?`) +- **migration_locations_get**: List migration locations (`instance?`) +- **private_ip_available_get**: List available private IPv4 addresses (`instance?`) +- **abuse_suspensions_get**: Get suspension details (`instance?`) +- **abuse_policy_get**: Get policy violations (`instance?`) +- **notification_preferences_get**: Get notification preferences (`instance?`) All MCP tools are safe, read-only operations that won't modify your VPS configuration or data. @@ -228,6 +243,8 @@ iso Manage ISO images for VPS boot reinstall Reinstall VPS operating system (WARNING: destroys all data) usage Display detailed VPS usage statistics audit Display audit log entries +abuse Display suspension details and policy violations +notifications Display KiwiVM notification preferences reset-password Reset the root password snapshot Manage VPS snapshots backup Manage VPS backups diff --git a/README.zh.md b/README.zh.md index 6e06f5a..0e9c05b 100644 --- a/README.zh.md +++ b/README.zh.md @@ -30,8 +30,11 @@ bwh node add production --api-key --veid # 基本操作 bwh info # 查看服务器详情 +bwh rate-limit # 检查 API 限制状态 bwh start/stop/restart # 电源管理 bwh usage --period 7d # 检查使用统计 +bwh abuse suspensions # 查看暂停详情 +bwh notifications list # 查看通知偏好 bwh snapshot create "备份名称" # 创建快照 bwh iso images # 列出可用 ISO 镜像 bwh iso mount ubuntu-20.04.iso # 挂载 ISO 用于救援/安装 @@ -98,12 +101,16 @@ backups, err := c.ListBackups(ctx) **服务器管理**: `GetServiceInfo`、`GetLiveServiceInfo`、`Start`、`Stop`、`Restart`、`Kill`、`SetHostname`、`ReinstallOS`、`ResetRootPassword`、`MountISO`、`UnmountISO` -**监控**: `GetRawUsageStats`、`GetBasicServiceInfo`、审计日志访问 +**监控**: `GetRawUsageStats`、`GetAuditLog`、`GetRateLimitStatus` **备份和恢复**: `CreateSnapshot`、`RestoreSnapshot`、`DeleteSnapshot`、备份管理 **迁移**: `GetMigrateLocations`、`StartMigration`（支持 `StartMigrationWithTimeout` 自定义超时） +**安全与 abuse**: `GetSuspensionDetails`、`GetPolicyViolations` + +**通知**: `GetNotificationPreferences` + **网络**: SSH 密钥管理、IP/反向 DNS 配置、IPv6 子网管理、私有 IPv4 管理 *完整 API 参考*: 查看 [pkg/client 文档](./pkg/client) 或运行 `go doc github.com/strahe/bwh/pkg/client` 获取所有可用方法。 @@ -187,7 +194,7 @@ claude mcp add bwh -- bwh mcp serve ### 配置说明 - **自定义配置**: 使用 `--config /path/to/config.yaml` 指定配置文件 -- **多实例**: 服务器自动使用配置中的默认实例 +- **多实例**: 服务器使用配置中的默认实例，或使用 `--instance ` 作为 MCP 会话默认实例。工具级 `instance` 参数可覆盖该默认值。 - **集成**: 添加到现有 MCP 配置文件中，不替换其他服务器 ### 可用工具 @@ -199,6 +206,14 @@ claude mcp add bwh -- bwh mcp serve - **backup_list**: 列出备份 (`instance?`, `os_contains?`, `since?`, `until?`, `sort_by?`, `order?`, `limit?`) - **vps_audit_get**: 获取审计日志 (`instance?`, `since?`, `until?`, `limit?`, `ip_contains?`, `type?`) - **iso_list**: 列出可用和已挂载的 ISO 镜像 (`instance?`) +- **ssh_keys_get**: 获取 SSH 公钥 (`instance?`, `full?`) +- **os_templates_get**: 列出系统模板 (`instance?`) +- **rate_limit_get**: 获取 API 限制状态 (`instance?`) +- **migration_locations_get**: 列出迁移位置 (`instance?`) +- **private_ip_available_get**: 列出可用私有 IPv4 地址 (`instance?`) +- **abuse_suspensions_get**: 获取暂停详情 (`instance?`) +- **abuse_policy_get**: 获取策略违规 (`instance?`) +- **notification_preferences_get**: 获取通知偏好 (`instance?`) 所有 MCP 工具都是安全的只读操作，不会修改您的 VPS 配置或数据。 @@ -228,6 +243,8 @@ iso 管理 VPS 启动用 ISO 镜像 reinstall 重装 VPS 操作系统（警告：摧毁所有数据） usage 显示详细 VPS 使用统计 audit 显示审计日志条目 +abuse 显示暂停详情与策略违规 +notifications 显示 KiwiVM 通知偏好 reset-password 重置 root 密码 snapshot 管理 VPS 快照 backup 管理 VPS 备份 diff --git a/cmd/bwh/abuse.go b/cmd/bwh/abuse.go new file mode 100644 index 0000000..ce6ca0a --- /dev/null +++ b/cmd/bwh/abuse.go @@ -0,0 +1,133 @@ +package main + +import ( + "context" + "fmt" + "sort" + "time" + + "github.com/strahe/bwh/pkg/client" + "github.com/urfave/cli/v3" +) + +var abuseCmd = &cli.Command{ + Name: "abuse", + Usage: "inspect abuse suspensions and policy violations", + Commands: []*cli.Command{ + abuseSuspensionsCmd, + abusePolicyCmd, + }, +} + +var abuseSuspensionsCmd = &cli.Command{ + Name: "suspensions", + Usage: "show service suspension details", + Action: func(ctx context.Context, cmd *cli.Command) error { + bwhClient, resolvedName, err := createBWHClient(cmd) + if err != nil { + return err + } + + fmt.Printf("Getting suspension details for instance: %s\n", resolvedName) + resp, err := bwhClient.GetSuspensionDetails(ctx) + if err != nil { + return fmt.Errorf("failed to get suspension details: %w", err) + } + + displaySuspensionDetails(resp) + return nil + }, +} + +var abusePolicyCmd = &cli.Command{ + Name: "policy", + Usage: "show active policy violations", + Action: func(ctx context.Context, cmd *cli.Command) error { + bwhClient, resolvedName, err := createBWHClient(cmd) + if err != nil { + return err + } + + fmt.Printf("Getting policy violations for instance: %s\n", resolvedName) + resp, err := bwhClient.GetPolicyViolations(ctx) + if err != nil { + return fmt.Errorf("failed to get policy violations: %w", err) + } + + displayPolicyViolations(resp) + return nil + }, +} + +func displaySuspensionDetails(resp *client.SuspensionDetailsResponse) { + fmt.Printf("\n🚫 SUSPENSION DETAILS\n") + fmt.Printf(" Suspensions (YTD): %d\n", resp.SuspensionCount) + fmt.Printf(" Abuse Points : %d / %d\n", resp.TotalAbusePoints, resp.MaxAbusePoints) + + if len(resp.Suspensions) == 0 { + fmt.Printf("\nNo active suspension issues found.\n") + return + } + + records := append([]client.SuspensionRecord{}, resp.Suspensions...) + sort.Slice(records, func(i, j int) bool { + return records[i].RecordID < records[j].RecordID + }) + + fmt.Printf("\nOutstanding Issues (%d):\n", len(records)) + for i, record := range records { + fmt.Printf("\n[%d] Case #%d\n", i+1, record.RecordID) + fmt.Printf(" Flag : %s\n", record.Flag) + fmt.Printf(" Soft Resolve: %s\n", yesNo(record.IsSoft == 1)) + fmt.Printf(" Abuse Points: %d\n", record.AbusePoints) + if record.EvidenceRecordID != 0 { + fmt.Printf(" Evidence ID : %d\n", record.EvidenceRecordID) + if text := resp.Evidence[fmt.Sprintf("%d", record.EvidenceRecordID)]; text != "" { + fmt.Printf(" Evidence : %s\n", summarizeText(text, 120)) + } + } + } +} + +func displayPolicyViolations(resp *client.PolicyViolationsResponse) { + fmt.Printf("\n⚠️ POLICY VIOLATIONS\n") + fmt.Printf(" Abuse Points: %d / %d\n", resp.TotalAbusePoints, resp.MaxAbusePoints) + + if len(resp.PolicyViolations) == 0 { + fmt.Printf("\nNo active policy violations found.\n") + return + } + + records := append([]client.PolicyViolationRecord{}, resp.PolicyViolations...) + sort.Slice(records, func(i, j int) bool { + return records[i].Timestamp > records[j].Timestamp + }) + + fmt.Printf("\nActive Violations (%d):\n", len(records)) + for i, record := range records { + fmt.Printf("\n[%d] Case #%d\n", i+1, record.RecordID) + fmt.Printf(" Flag : %s\n", record.Flag) + fmt.Printf(" Soft Resolve: %s\n", yesNo(record.IsSoft == 1)) + fmt.Printf(" Abuse Points: %d\n", record.AbusePoints) + if record.Timestamp > 0 { + fmt.Printf(" Created : %s\n", time.Unix(record.Timestamp, 0).Format("2006-01-02 15:04:05")) + } + if record.SuspendAt > 0 { + fmt.Printf(" Suspend At : %s\n", time.Unix(record.SuspendAt, 0).Format("2006-01-02 15:04:05")) + } + if record.EvidenceData != "" { + fmt.Printf(" Evidence : %s\n", summarizeText(record.EvidenceData, 120)) + } + } +} + +func summarizeText(s string, maxLen int) string { + if maxLen <= 0 { + return "..." + } + runes := []rune(s) + if len(runes) <= maxLen { + return s + } + return string(runes[:maxLen]) + "..." +} diff --git a/cmd/bwh/info.go b/cmd/bwh/info.go index d1cba2a..5e0bbc6 100644 --- a/cmd/bwh/info.go +++ b/cmd/bwh/info.go @@ -3,6 +3,7 @@ package main import ( "context" "fmt" + "sort" "strings" "time" @@ -176,8 +177,9 @@ func displayDetailedInfo(info *client.LiveServiceInfo, instanceName string) { } if len(info.IPNullroutes) > 0 { + nullrouteIPs := sortedNullrouteIPs(info.IPNullroutes) fmt.Printf(" ⚠️ DDoS Protection : %d IP(s) currently null-routed\n", len(info.IPNullroutes)) - fmt.Printf(" Null-routed IPs : %s\n", strings.Join(info.IPNullroutes, ", ")) + fmt.Printf(" Null-routed IPs : %s\n", strings.Join(nullrouteIPs, ", ")) } if info.IPv6SitTunnelEndpoint != "" { @@ -444,6 +446,15 @@ func displaySecurityInfo(info *client.ServiceInfo) { } } +func sortedNullrouteIPs(nullroutes client.IPNullroutes) []string { + ips := make([]string, 0, len(nullroutes)) + for ip := range nullroutes { + ips = append(ips, ip) + } + sort.Strings(ips) + return ips +} + // formatBool converts boolean to readable yes/no func formatBool(b bool) string { if b { diff --git a/cmd/bwh/main.go b/cmd/bwh/main.go index 2c5ca7c..6c9c36f 100644 --- a/cmd/bwh/main.go +++ b/cmd/bwh/main.go @@ -52,6 +52,8 @@ func main() { reinstallCmd, usageCmd, auditCmd, + abuseCmd, + notificationsCmd, resetPasswordCmd, snapshotCmd, backupCmd, diff --git a/cmd/bwh/notifications.go b/cmd/bwh/notifications.go new file mode 100644 index 0000000..3aa025a --- /dev/null +++ b/cmd/bwh/notifications.go @@ -0,0 +1,89 @@ +package main + +import ( + "context" + "fmt" + "sort" + "time" + + "github.com/strahe/bwh/pkg/client" + "github.com/urfave/cli/v3" +) + +var notificationsCmd = &cli.Command{ + Name: "notifications", + Usage: "inspect KiwiVM notification preferences", + Commands: []*cli.Command{ + notificationsListCmd, + }, +} + +var notificationsListCmd = &cli.Command{ + Name: "list", + Usage: "list KiwiVM notification preferences", + Action: func(ctx context.Context, cmd *cli.Command) error { + bwhClient, resolvedName, err := createBWHClient(cmd) + if err != nil { + return err + } + + fmt.Printf("Getting notification preferences for instance: %s\n", resolvedName) + resp, err := bwhClient.GetNotificationPreferences(ctx) + if err != nil { + return fmt.Errorf("failed to get notification preferences: %w", err) + } + + displayNotificationPreferences(resp) + return nil + }, +} + +func displayNotificationPreferences(resp *client.NotificationPreferencesResponse) { + fmt.Printf("\n📧 NOTIFICATION PREFERENCES\n") + if resp.NotificationEmail != "" { + fmt.Printf(" Email: %s\n", resp.NotificationEmail) + } + + if len(resp.EmailPreferences) == 0 { + fmt.Printf("\nNo notification preferences found.\n") + return + } + + categories := make([]string, 0, len(resp.EmailPreferences)) + for category := range resp.EmailPreferences { + categories = append(categories, category) + } + sort.Strings(categories) + + for _, category := range categories { + prefs := resp.EmailPreferences[category] + ids := make([]string, 0, len(prefs)) + for id := range prefs { + ids = append(ids, id) + } + sort.Strings(ids) + + fmt.Printf("\n%s\n", category) + for _, id := range ids { + pref := prefs[id] + fmt.Printf(" • %s\n", id) + fmt.Printf(" Status : %s\n", enabledStatus(pref.IsEnabled)) + if pref.FriendlyDescription != "" { + fmt.Printf(" Description: %s\n", pref.FriendlyDescription) + } + if pref.ChangedTimestamp > 0 { + fmt.Printf(" Updated : %s\n", time.Unix(pref.ChangedTimestamp, 0).Format("2006-01-02 15:04:05")) + } + if pref.SValue != "" { + fmt.Printf(" Value : %s\n", pref.SValue) + } + } + } +} + +func enabledStatus(value int) string { + if value == 1 { + return "✅ Enabled" + } + return "❌ Disabled" +} diff --git a/cmd/bwh/read_only_display_test.go b/cmd/bwh/read_only_display_test.go new file mode 100644 index 0000000..4fb93a2 --- /dev/null +++ b/cmd/bwh/read_only_display_test.go @@ -0,0 +1,187 @@ +package main + +import ( + "bytes" + "io" + "os" + "strings" + "testing" + + "github.com/strahe/bwh/pkg/client" +) + +func captureStdout(t *testing.T, fn func()) string { + t.Helper() + + old := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("failed to create stdout pipe: %v", err) + } + os.Stdout = w + + fn() + + if err := w.Close(); err != nil { + t.Fatalf("failed to close stdout writer: %v", err) + } + os.Stdout = old + + var buf bytes.Buffer + if _, err := io.Copy(&buf, r); err != nil { + t.Fatalf("failed to read stdout: %v", err) + } + return buf.String() +} + +func TestDisplaySuspensionDetails(t *testing.T) { + t.Run("empty", func(t *testing.T) { + out := captureStdout(t, func() { + displaySuspensionDetails(&client.SuspensionDetailsResponse{ + SuspensionCount: 0, + TotalAbusePoints: 0, + MaxAbusePoints: 60, + }) + }) + + if !strings.Contains(out, "No active suspension issues found.") { + t.Fatalf("output = %q", out) + } + }) + + t.Run("with records", func(t *testing.T) { + out := captureStdout(t, func() { + displaySuspensionDetails(&client.SuspensionDetailsResponse{ + SuspensionCount: 1, + TotalAbusePoints: 20, + MaxAbusePoints: 60, + Suspensions: []client.SuspensionRecord{ + { + RecordID: 123, + Flag: "spam", + IsSoft: 1, + EvidenceRecordID: 456, + AbusePoints: 20, + }, + }, + Evidence: map[string]string{"456": strings.Repeat("x", 140)}, + }) + }) + + for _, want := range []string{"Outstanding Issues (1):", "Case #123", "Soft Resolve:", strings.Repeat("x", 120) + "..."} { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } + }) +} + +func TestDisplayPolicyViolations(t *testing.T) { + t.Run("empty", func(t *testing.T) { + out := captureStdout(t, func() { + displayPolicyViolations(&client.PolicyViolationsResponse{ + TotalAbusePoints: 0, + MaxAbusePoints: 60, + }) + }) + + if !strings.Contains(out, "No active policy violations found.") { + t.Fatalf("output = %q", out) + } + }) + + t.Run("with records", func(t *testing.T) { + out := captureStdout(t, func() { + displayPolicyViolations(&client.PolicyViolationsResponse{ + TotalAbusePoints: 10, + MaxAbusePoints: 60, + PolicyViolations: []client.PolicyViolationRecord{ + { + RecordID: 789, + Timestamp: 1710000000, + SuspendAt: 1710003600, + Flag: "policy", + AbusePoints: 10, + EvidenceData: "sample policy evidence", + }, + }, + }) + }) + + for _, want := range []string{"Active Violations (1):", "Case #789", "sample policy evidence"} { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } + }) +} + +func TestDisplayNotificationPreferences(t *testing.T) { + t.Run("empty", func(t *testing.T) { + out := captureStdout(t, func() { + displayNotificationPreferences(&client.NotificationPreferencesResponse{}) + }) + + if !strings.Contains(out, "No notification preferences found.") { + t.Fatalf("output = %q", out) + } + }) + + t.Run("with records", func(t *testing.T) { + out := captureStdout(t, func() { + displayNotificationPreferences(&client.NotificationPreferencesResponse{ + NotificationEmail: "user@example.com", + EmailPreferences: map[string]map[string]client.NotificationPreference{ + "service": { + "maintenance": { + FriendlyDescription: "Maintenance notices", + IsEnabled: 1, + ChangedTimestamp: 1710000000, + SValue: "daily", + }, + }, + }, + }) + }) + + for _, want := range []string{"Email: user@example.com", "maintenance", "Maintenance notices", "Value : daily"} { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } + }) +} + +func TestDisplayDetailedInfoNullroutes(t *testing.T) { + out := captureStdout(t, func() { + displayDetailedInfo(&client.LiveServiceInfo{ + ServiceInfo: client.ServiceInfo{ + Hostname: "test-host", + VMType: "kvm", + Plan: "test-plan", + OS: "debian", + IPAddresses: []string{"198.51.100.10"}, + PlanMaxIPv6s: 1, + MaxAbusePoints: 60, + IPNullroutes: client.IPNullroutes{ + "192.0.2.20": {}, + "192.0.2.10": {}, + }, + }, + }, "test") + }) + + if !strings.Contains(out, "DDoS Protection : 2 IP(s) currently null-routed") { + t.Fatalf("output = %q", out) + } + if !strings.Contains(out, "Null-routed IPs : 192.0.2.10, 192.0.2.20") { + t.Fatalf("output = %q", out) + } +} + +func TestSummarizeTextUsesRunes(t *testing.T) { + got := summarizeText("证据内容abc", 4) + if got != "证据内容..." { + t.Fatalf("summarizeText() = %q", got) + } +} diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index 42a498d..db398ba 100644 --- a/internal/mcpserver/server.go +++ b/internal/mcpserver/server.go @@ -24,7 +24,7 @@ func RunMCPStdioServer(ctx context.Context, configPath, instanceName string) err } // Resolve once here only for connectivity check (uses provided instanceName or config default) - instForCheck, _, err := manager.ResolveInstance(instanceName) + instForCheck, resolvedInstanceName, err := manager.ResolveInstance(instanceName) if err != nil { return fmt.Errorf("failed to resolve instance: %w", err) } @@ -50,7 +50,7 @@ func RunMCPStdioServer(ctx context.Context, configPath, instanceName string) err ) // Register read-only tools - registerReadOnlyTools(s, manager) + registerReadOnlyTools(s, manager, resolvedInstanceName) // Register simple resources registerResources(s, manager) @@ -59,8 +59,163 @@ func RunMCPStdioServer(ctx context.Context, configPath, instanceName string) err return server.ServeStdio(s) } +func resolveClient(manager *config.Manager, requested, defaultInstance string) (*client.Client, string, error) { + target := strings.TrimSpace(requested) + if target == "" { + target = defaultInstance + } + + inst, resolved, err := manager.ResolveInstance(target) + if err != nil { + return nil, "", err + } + + c := client.NewClient(inst.APIKey, inst.VeID) + if inst.Endpoint != "" { + c.SetBaseURL(inst.Endpoint) + } + return c, resolved, nil +} + +func callReadOnlyTool[T any]( + ctx context.Context, + manager *config.Manager, + requested string, + defaultInstance string, + action string, + call func(context.Context, *client.Client) (T, error), + payload func(string, T) map[string]any, +) (*mcp.CallToolResult, error) { + c, resolved, err := resolveClient(manager, requested, defaultInstance) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil + } + + resp, err := call(ctx, c) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("%s failed: %v", action, err)), nil + } + + return mcp.NewToolResultStructuredOnly(payload(resolved, resp)), nil +} + +func sshKeysPayload(resolved string, resp *client.SshKeysResponse, full bool) map[string]any { + keys := map[string][]string{ + "veid": resp.GetShortenedSshKeysVeidSlice(), + "user": resp.GetShortenedSshKeysUserSlice(), + "preferred": resp.GetShortenedSshKeysPreferredSlice(), + } + if full { + keys = map[string][]string{ + "veid": resp.GetSshKeysVeidSlice(), + "user": resp.GetSshKeysUserSlice(), + "preferred": resp.GetSshKeysPreferredSlice(), + } + } + + return map[string]any{ + "instance": resolved, + "full": full, + "keys": keys, + "totals": map[string]int{ + "veid": len(keys["veid"]), + "user": len(keys["user"]), + "preferred": len(keys["preferred"]), + }, + } +} + +func availableOSPayload(resolved string, resp *client.AvailableOSResponse) map[string]any { + return map[string]any{ + "instance": resolved, + "installed": resp.Installed, + "templates": resp.Templates, + "total_templates": len(resp.Templates), + } +} + +func rateLimitPayload(resolved string, resp *client.RateLimitStatus) map[string]any { + return map[string]any{ + "instance": resolved, + "remaining_points_15min": resp.RemainingPoints15Min, + "remaining_points_24h": resp.RemainingPoints24H, + } +} + +func migrationLocationsPayload(resolved string, resp *client.MigrateLocationsResponse) map[string]any { + type locationInfo struct { + ID string `json:"id"` + Description string `json:"description,omitempty"` + DataTransferMultiple int `json:"data_transfer_multiplier,omitempty"` + IsCurrent bool `json:"is_current"` + } + + locations := make([]locationInfo, 0, len(resp.Locations)) + for _, id := range resp.Locations { + locations = append(locations, locationInfo{ + ID: id, + Description: resp.Descriptions[id], + DataTransferMultiple: resp.DataTransferMultipliers[id], + IsCurrent: id == resp.CurrentLocation, + }) + } + + return map[string]any{ + "instance": resolved, + "current_location": resp.CurrentLocation, + "locations": locations, + "total": len(locations), + } +} + +func privateIPAvailablePayload(resolved string, resp *client.PrivateIPAvailableResponse) map[string]any { + return map[string]any{ + "instance": resolved, + "available_ips": resp.AvailableIPs, + "total": len(resp.AvailableIPs), + } +} + +func suspensionDetailsPayload(resolved string, resp *client.SuspensionDetailsResponse) map[string]any { + return map[string]any{ + "instance": resolved, + "suspension_count": resp.SuspensionCount, + "total_abuse_points": resp.TotalAbusePoints, + "max_abuse_points": resp.MaxAbusePoints, + "suspensions": resp.Suspensions, + "evidence": resp.Evidence, + } +} + +func policyViolationsPayload(resolved string, resp *client.PolicyViolationsResponse) map[string]any { + return map[string]any{ + "instance": resolved, + "total_abuse_points": resp.TotalAbusePoints, + "max_abuse_points": resp.MaxAbusePoints, + "policy_violations": resp.PolicyViolations, + } +} + +func notificationPreferencesPayload(resolved string, resp *client.NotificationPreferencesResponse) map[string]any { + return map[string]any{ + "instance": resolved, + "notification_email": resp.NotificationEmail, + "email_preferences": resp.EmailPreferences, + "total_categories": len(resp.EmailPreferences), + "total_preference_ids": countNotificationPreferences(resp.EmailPreferences), + } +} + +func countNotificationPreferences(prefs map[string]map[string]client.NotificationPreference) int { + total := 0 + for _, category := range prefs { + total += len(category) + } + return total +} + // registerReadOnlyTools wires read-only tool handlers backed by pkg/client -func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { +func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager, defaultInstance string) { // vps_info_get s.AddTool( mcp.NewTool( @@ -79,14 +234,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { requested := req.GetString("instance", "") compact := req.GetBool("compact", false) live := req.GetBool("live", true) - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } if live { info, err := c.GetLiveServiceInfo(ctx) if err != nil { @@ -157,14 +308,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { daysArg := req.GetInt("days", 0) groupBy := req.GetString("group_by", "day") - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } stats, err := c.GetRawUsageStats(ctx) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("get usage failed: %v", err)), nil @@ -379,14 +526,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { order := req.GetString("order", "asc") limit := req.GetInt("limit", 0) - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } list, err := c.ListSnapshots(ctx) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("list snapshots failed: %v", err)), nil @@ -465,14 +608,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { order := req.GetString("order", "desc") limit := req.GetInt("limit", 0) - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } resp, err := c.ListBackups(ctx) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("list backups failed: %v", err)), nil @@ -554,14 +693,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { ipContains := strings.TrimSpace(req.GetString("ip_contains", "")) typeFilter := req.GetInt("type", -1) - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } logResp, err := c.GetAuditLog(ctx) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("get audit log failed: %v", err)), nil @@ -627,14 +762,10 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { requested := req.GetString("instance", "") - inst, resolved, err := manager.ResolveInstance(requested) + c, resolved, err := resolveClient(manager, requested, defaultInstance) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("resolve instance failed: %v", err)), nil } - c := client.NewClient(inst.APIKey, inst.VeID) - if inst.Endpoint != "" { - c.SetBaseURL(inst.Endpoint) - } serviceInfo, err := c.GetServiceInfo(ctx) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("get service info failed: %v", err)), nil @@ -659,6 +790,208 @@ func registerReadOnlyTools(s *server.MCPServer, manager *config.Manager) { }, ) + // ssh_keys_get + s.AddTool( + mcp.NewTool( + "ssh_keys_get", + mcp.WithDescription("Get SSH public keys for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + mcp.WithBoolean("full", mcp.DefaultBool(false), mcp.Description("Return full public keys instead of shortened keys")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + full := req.GetBool("full", false) + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get SSH keys", + func(ctx context.Context, c *client.Client) (*client.SshKeysResponse, error) { + return c.GetSshKeys(ctx) + }, + func(resolved string, resp *client.SshKeysResponse) map[string]any { + return sshKeysPayload(resolved, resp, full) + }, + ) + }, + ) + + // os_templates_get + s.AddTool( + mcp.NewTool( + "os_templates_get", + mcp.WithDescription("List OS templates available for reinstall on BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get OS templates", + func(ctx context.Context, c *client.Client) (*client.AvailableOSResponse, error) { + return c.GetAvailableOS(ctx) + }, + func(resolved string, resp *client.AvailableOSResponse) map[string]any { + return availableOSPayload(resolved, resp) + }, + ) + }, + ) + + // rate_limit_get + s.AddTool( + mcp.NewTool( + "rate_limit_get", + mcp.WithDescription("Get current API rate limit status for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + func(resolved string, resp *client.RateLimitStatus) map[string]any { + return rateLimitPayload(resolved, resp) + }, + ) + }, + ) + + // migration_locations_get + s.AddTool( + mcp.NewTool( + "migration_locations_get", + mcp.WithDescription("List VPS migration locations for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get migration locations", + func(ctx context.Context, c *client.Client) (*client.MigrateLocationsResponse, error) { + return c.GetMigrateLocations(ctx) + }, + func(resolved string, resp *client.MigrateLocationsResponse) map[string]any { + return migrationLocationsPayload(resolved, resp) + }, + ) + }, + ) + + // private_ip_available_get + s.AddTool( + mcp.NewTool( + "private_ip_available_get", + mcp.WithDescription("List available private IPv4 addresses for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get available private IPs", + func(ctx context.Context, c *client.Client) (*client.PrivateIPAvailableResponse, error) { + return c.GetAvailablePrivateIPs(ctx) + }, + func(resolved string, resp *client.PrivateIPAvailableResponse) map[string]any { + return privateIPAvailablePayload(resolved, resp) + }, + ) + }, + ) + + // abuse_suspensions_get + s.AddTool( + mcp.NewTool( + "abuse_suspensions_get", + mcp.WithDescription("Get service suspension details and abuse evidence for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get suspension details", + func(ctx context.Context, c *client.Client) (*client.SuspensionDetailsResponse, error) { + return c.GetSuspensionDetails(ctx) + }, + func(resolved string, resp *client.SuspensionDetailsResponse) map[string]any { + return suspensionDetailsPayload(resolved, resp) + }, + ) + }, + ) + + // abuse_policy_get + s.AddTool( + mcp.NewTool( + "abuse_policy_get", + mcp.WithDescription("Get active policy violations for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get policy violations", + func(ctx context.Context, c *client.Client) (*client.PolicyViolationsResponse, error) { + return c.GetPolicyViolations(ctx) + }, + func(resolved string, resp *client.PolicyViolationsResponse) map[string]any { + return policyViolationsPayload(resolved, resp) + }, + ) + }, + ) + + // notification_preferences_get + s.AddTool( + mcp.NewTool( + "notification_preferences_get", + mcp.WithDescription("Get KiwiVM notification preferences for BWH/BandwagonHost/搬瓦工/瓦工"), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithOpenWorldHintAnnotation(true), + mcp.WithString("instance", mcp.Description("Target instance name; defaults to config default")), + ), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + requested := req.GetString("instance", "") + + return callReadOnlyTool(ctx, manager, requested, defaultInstance, "get notification preferences", + func(ctx context.Context, c *client.Client) (*client.NotificationPreferencesResponse, error) { + return c.GetNotificationPreferences(ctx) + }, + func(resolved string, resp *client.NotificationPreferencesResponse) map[string]any { + return notificationPreferencesPayload(resolved, resp) + }, + ) + }, + ) + // instance_list s.AddTool( mcp.NewTool( diff --git a/internal/mcpserver/server_test.go b/internal/mcpserver/server_test.go new file mode 100644 index 0000000..5df44cb --- /dev/null +++ b/internal/mcpserver/server_test.go @@ -0,0 +1,313 @@ +package mcpserver + +import ( + "context" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + + "github.com/strahe/bwh/internal/config" + "github.com/strahe/bwh/pkg/client" +) + +func newMCPTestManager(t *testing.T, endpoint string) *config.Manager { + t.Helper() + + manager, err := config.NewManager(filepath.Join(t.TempDir(), "config.yaml")) + if err != nil { + t.Fatalf("NewManager() error = %v", err) + } + + err = manager.AddInstance("default", &config.Instance{ + APIKey: "test-api-key-123456789", + VeID: "123456", + Endpoint: endpoint, + }, true) + if err != nil { + t.Fatalf("AddInstance() error = %v", err) + } + return manager +} + +func TestMCPReadOnlyPayloads(t *testing.T) { + t.Run("ssh keys default shortened", func(t *testing.T) { + payload := sshKeysPayload("default", &client.SshKeysResponse{ + SshKeysVeid: "ssh-rsa full-vm", + SshKeysUser: "ssh-rsa full-user", + SshKeysPreferred: "ssh-rsa full-preferred", + ShortenedSshKeysVeid: "ssh-rsa short-vm", + ShortenedSshKeysUser: "ssh-rsa short-user", + ShortenedSshKeysPreferred: "ssh-rsa short-preferred", + }, false) + + keys := payload["keys"].(map[string][]string) + if keys["veid"][0] != "ssh-rsa short-vm" { + t.Fatalf("veid key = %q", keys["veid"][0]) + } + if payload["full"].(bool) { + t.Fatal("full = true, want false") + } + }) + + t.Run("ssh keys full", func(t *testing.T) { + payload := sshKeysPayload("default", &client.SshKeysResponse{ + SshKeysVeid: "ssh-rsa full-vm", + SshKeysUser: "ssh-rsa full-user", + SshKeysPreferred: "ssh-rsa full-preferred", + ShortenedSshKeysVeid: "ssh-rsa short-vm", + }, true) + + keys := payload["keys"].(map[string][]string) + if keys["veid"][0] != "ssh-rsa full-vm" { + t.Fatalf("veid key = %q", keys["veid"][0]) + } + }) + + t.Run("os templates", func(t *testing.T) { + payload := availableOSPayload("default", &client.AvailableOSResponse{ + Installed: "debian-12-x86_64", + Templates: []string{ + "debian-12-x86_64", + "ubuntu-24.04-x86_64", + }, + }) + + if payload["total_templates"] != 2 { + t.Fatalf("total_templates = %v", payload["total_templates"]) + } + }) + + t.Run("rate limit", func(t *testing.T) { + payload := rateLimitPayload("default", &client.RateLimitStatus{ + RemainingPoints15Min: 997, + RemainingPoints24H: 19852, + }) + + if payload["remaining_points_15min"] != 997 { + t.Fatalf("remaining_points_15min = %v", payload["remaining_points_15min"]) + } + }) + + t.Run("migration locations", func(t *testing.T) { + payload := migrationLocationsPayload("default", &client.MigrateLocationsResponse{ + CurrentLocation: "usca_2", + Locations: []string{"usca_2", "usny_6"}, + Descriptions: map[string]string{ + "usca_2": "US: California", + "usny_6": "US: New York", + }, + DataTransferMultipliers: map[string]int{ + "usca_2": 1, + "usny_6": 2, + }, + }) + + if payload["total"] != 2 { + t.Fatalf("total = %v", payload["total"]) + } + }) + + t.Run("private ip", func(t *testing.T) { + payload := privateIPAvailablePayload("default", &client.PrivateIPAvailableResponse{ + AvailableIPs: []string{"10.0.0.10", "10.0.0.11"}, + }) + + if payload["total"] != 2 { + t.Fatalf("total = %v", payload["total"]) + } + }) + + t.Run("suspensions", func(t *testing.T) { + payload := suspensionDetailsPayload("default", &client.SuspensionDetailsResponse{ + SuspensionCount: 1, + TotalAbusePoints: 20, + MaxAbusePoints: 60, + Suspensions: []client.SuspensionRecord{ + {RecordID: 123, Flag: "spam", AbusePoints: 20}, + }, + Evidence: map[string]string{"456": "sample evidence"}, + }) + + if payload["suspension_count"] != 1 { + t.Fatalf("suspension_count = %v", payload["suspension_count"]) + } + }) + + t.Run("policy violations", func(t *testing.T) { + payload := policyViolationsPayload("default", &client.PolicyViolationsResponse{ + TotalAbusePoints: 10, + MaxAbusePoints: 60, + PolicyViolations: []client.PolicyViolationRecord{ + {RecordID: 789, Flag: "policy", AbusePoints: 10}, + }, + }) + + if payload["total_abuse_points"] != 10 { + t.Fatalf("total_abuse_points = %v", payload["total_abuse_points"]) + } + }) + + t.Run("notifications", func(t *testing.T) { + payload := notificationPreferencesPayload("default", &client.NotificationPreferencesResponse{ + NotificationEmail: "user@example.com", + EmailPreferences: map[string]map[string]client.NotificationPreference{ + "service": { + "maintenance": {IsEnabled: 1}, + "billing": {IsEnabled: 0}, + }, + }, + }) + + if payload["total_categories"] != 1 { + t.Fatalf("total_categories = %v", payload["total_categories"]) + } + if payload["total_preference_ids"] != 2 { + t.Fatalf("total_preference_ids = %v", payload["total_preference_ids"]) + } + }) +} + +func TestCallReadOnlyTool(t *testing.T) { + t.Run("success", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := strings.TrimPrefix(r.URL.Path, "/"); got != "getRateLimitStatus" { + t.Fatalf("endpoint = %s, want getRateLimitStatus", got) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":0,"remaining_points_15min":997,"remaining_points_24h":19852}`)) + })) + defer server.Close() + + manager := newMCPTestManager(t, server.URL) + result, err := callReadOnlyTool(context.Background(), manager, "", "default", "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + rateLimitPayload, + ) + if err != nil { + t.Fatalf("callReadOnlyTool() error = %v", err) + } + if result.IsError { + t.Fatalf("result.IsError = true") + } + + payload := result.StructuredContent.(map[string]any) + if payload["instance"] != "default" { + t.Fatalf("instance = %v", payload["instance"]) + } + }) + + t.Run("resolve error", func(t *testing.T) { + manager, err := config.NewManager(filepath.Join(t.TempDir(), "config.yaml")) + if err != nil { + t.Fatalf("NewManager() error = %v", err) + } + + result, err := callReadOnlyTool(context.Background(), manager, "", "default", "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + rateLimitPayload, + ) + if err != nil { + t.Fatalf("callReadOnlyTool() error = %v", err) + } + if !result.IsError { + t.Fatalf("result.IsError = false") + } + }) + + t.Run("api error", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":700005,"message":"Authentication failure"}`)) + })) + defer server.Close() + + manager := newMCPTestManager(t, server.URL) + result, err := callReadOnlyTool(context.Background(), manager, "", "default", "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + rateLimitPayload, + ) + if err != nil { + t.Fatalf("callReadOnlyTool() error = %v", err) + } + if !result.IsError { + t.Fatalf("result.IsError = false") + } + }) +} + +func TestCallReadOnlyToolUsesServerDefaultInstance(t *testing.T) { + seenVEIDs := []string{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seenVEIDs = append(seenVEIDs, r.URL.Query().Get("veid")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":0,"remaining_points_15min":997,"remaining_points_24h":19852}`)) + })) + defer server.Close() + + manager, err := config.NewManager(filepath.Join(t.TempDir(), "config.yaml")) + if err != nil { + t.Fatalf("NewManager() error = %v", err) + } + if err := manager.AddInstance("primary", &config.Instance{ + APIKey: "primary-api-key-123456789", + VeID: "111111", + Endpoint: server.URL, + }, true); err != nil { + t.Fatalf("AddInstance(primary) error = %v", err) + } + if err := manager.AddInstance("secondary", &config.Instance{ + APIKey: "secondary-api-key-123456789", + VeID: "222222", + Endpoint: server.URL, + }, false); err != nil { + t.Fatalf("AddInstance(secondary) error = %v", err) + } + + result, err := callReadOnlyTool(context.Background(), manager, "", "secondary", "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + rateLimitPayload, + ) + if err != nil { + t.Fatalf("callReadOnlyTool() error = %v", err) + } + if result.IsError { + t.Fatalf("result.IsError = true") + } + payload := result.StructuredContent.(map[string]any) + if payload["instance"] != "secondary" { + t.Fatalf("instance = %v, want secondary", payload["instance"]) + } + if len(seenVEIDs) != 1 || seenVEIDs[0] != "222222" { + t.Fatalf("veids = %v, want [222222]", seenVEIDs) + } + + result, err = callReadOnlyTool(context.Background(), manager, "primary", "secondary", "get rate limit", + func(ctx context.Context, c *client.Client) (*client.RateLimitStatus, error) { + return c.GetRateLimitStatus(ctx) + }, + rateLimitPayload, + ) + if err != nil { + t.Fatalf("callReadOnlyTool() with explicit instance error = %v", err) + } + if result.IsError { + t.Fatalf("result.IsError = true for explicit instance") + } + payload = result.StructuredContent.(map[string]any) + if payload["instance"] != "primary" { + t.Fatalf("instance = %v, want primary", payload["instance"]) + } + if len(seenVEIDs) != 2 || seenVEIDs[1] != "111111" { + t.Fatalf("veids = %v, want second call to use 111111", seenVEIDs) + } +} diff --git a/pkg/client/client.go b/pkg/client/client.go index c3cdedf..e75f4f8 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -325,6 +325,36 @@ func (c *Client) GetRateLimitStatus(ctx context.Context) (*RateLimitStatus, erro return wrapErrorWithBase(&resp, resp.BaseResponse) } +// GetSuspensionDetails gets service suspension details and abuse evidence. +func (c *Client) GetSuspensionDetails(ctx context.Context) (*SuspensionDetailsResponse, error) { + var resp SuspensionDetailsResponse + if err := c.doRequest(ctx, "getSuspensionDetails", nil, &resp); err != nil { + return nil, err + } + + return wrapErrorWithBase(&resp, resp.BaseResponse) +} + +// GetPolicyViolations gets active policy violations. +func (c *Client) GetPolicyViolations(ctx context.Context) (*PolicyViolationsResponse, error) { + var resp PolicyViolationsResponse + if err := c.doRequest(ctx, "getPolicyViolations", nil, &resp); err != nil { + return nil, err + } + + return wrapErrorWithBase(&resp, resp.BaseResponse) +} + +// GetNotificationPreferences gets KiwiVM notification settings and their state. +func (c *Client) GetNotificationPreferences(ctx context.Context) (*NotificationPreferencesResponse, error) { + var resp NotificationPreferencesResponse + if err := c.doRequest(ctx, "kiwivm/getNotificationPreferences", nil, &resp); err != nil { + return nil, err + } + + return wrapErrorWithBase(&resp, resp.BaseResponse) +} + // GetSshKeys gets SSH keys from both Hypervisor Vault and Billing Portal func (c *Client) GetSshKeys(ctx context.Context) (*SshKeysResponse, error) { var resp SshKeysResponse diff --git a/pkg/client/read_only_methods_test.go b/pkg/client/read_only_methods_test.go new file mode 100644 index 0000000..db51730 --- /dev/null +++ b/pkg/client/read_only_methods_test.go @@ -0,0 +1,277 @@ +package client + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newReadOnlyMockServer(t *testing.T, responses map[string]string, seen *[]string) *httptest.Server { + t.Helper() + + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + path := strings.TrimPrefix(r.URL.Path, "/") + *seen = append(*seen, path) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + if body, ok := responses[path]; ok { + if _, err := w.Write([]byte(body)); err != nil { + t.Fatalf("failed to write mock response: %v", err) + } + return + } + if _, err := w.Write([]byte(`{"error":404,"message":"Endpoint not found"}`)); err != nil { + t.Fatalf("failed to write mock response: %v", err) + } + })) +} + +func TestClient_GetSuspensionDetails_Mock(t *testing.T) { + tests := []struct { + name string + body string + wantRecords int + wantEvidence string + }{ + { + name: "no records", + body: `{ + "error": 0, + "suspension_count": 0, + "total_abuse_points": 0, + "max_abuse_points": 60 + }`, + wantRecords: 0, + }, + { + name: "with records", + body: `{ + "error": 0, + "suspension_count": 1, + "total_abuse_points": 20, + "max_abuse_points": 60, + "suspensions": [ + { + "record_id": 123, + "flag": "spam", + "is_soft": 1, + "evidence_record_id": 456, + "abuse_points": 20 + } + ], + "evidence": { + "456": "sample evidence" + } + }`, + wantRecords: 1, + wantEvidence: "sample evidence", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + seen := []string{} + server := newReadOnlyMockServer(t, map[string]string{ + "getSuspensionDetails": tt.body, + }, &seen) + defer server.Close() + + c := NewClient("valid_key", "123456") + c.SetBaseURL(server.URL) + + resp, err := c.GetSuspensionDetails(context.Background()) + if err != nil { + t.Fatalf("GetSuspensionDetails() error = %v", err) + } + + if len(seen) != 1 || seen[0] != "getSuspensionDetails" { + t.Fatalf("endpoint = %v, want [getSuspensionDetails]", seen) + } + if len(resp.Suspensions) != tt.wantRecords { + t.Fatalf("suspensions length = %d, want %d", len(resp.Suspensions), tt.wantRecords) + } + if tt.wantEvidence != "" && resp.Evidence["456"] != tt.wantEvidence { + t.Errorf("evidence = %q, want %q", resp.Evidence["456"], tt.wantEvidence) + } + }) + } +} + +func TestClient_GetPolicyViolations_Mock(t *testing.T) { + tests := []struct { + name string + body string + wantRecords int + }{ + { + name: "no records", + body: `{ + "error": 0, + "total_abuse_points": 0, + "max_abuse_points": 60 + }`, + wantRecords: 0, + }, + { + name: "with records", + body: `{ + "error": 0, + "total_abuse_points": 10, + "max_abuse_points": 60, + "policy_violations": [ + { + "record_id": 789, + "timestamp": 1710000000, + "suspend_at": 1710003600, + "flag": "policy", + "is_soft": 0, + "abuse_points": 10, + "evidence_data": "sample policy evidence" + } + ] + }`, + wantRecords: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + seen := []string{} + server := newReadOnlyMockServer(t, map[string]string{ + "getPolicyViolations": tt.body, + }, &seen) + defer server.Close() + + c := NewClient("valid_key", "123456") + c.SetBaseURL(server.URL) + + resp, err := c.GetPolicyViolations(context.Background()) + if err != nil { + t.Fatalf("GetPolicyViolations() error = %v", err) + } + + if len(seen) != 1 || seen[0] != "getPolicyViolations" { + t.Fatalf("endpoint = %v, want [getPolicyViolations]", seen) + } + if len(resp.PolicyViolations) != tt.wantRecords { + t.Fatalf("policy_violations length = %d, want %d", len(resp.PolicyViolations), tt.wantRecords) + } + }) + } +} + +func TestClient_GetNotificationPreferences_Mock(t *testing.T) { + seen := []string{} + server := newReadOnlyMockServer(t, map[string]string{ + "kiwivm/getNotificationPreferences": `{ + "error": 0, + "notificationEmail": "user@example.com", + "email_preferences": { + "service": { + "maintenance": { + "friendly_description": "Maintenance notices", + "is_enabled": 1, + "changed_timestamp": 1710000000, + "s_value": "daily" + } + } + } + }`, + }, &seen) + defer server.Close() + + c := NewClient("valid_key", "123456") + c.SetBaseURL(server.URL) + + resp, err := c.GetNotificationPreferences(context.Background()) + if err != nil { + t.Fatalf("GetNotificationPreferences() error = %v", err) + } + + if len(seen) != 1 || seen[0] != "kiwivm/getNotificationPreferences" { + t.Fatalf("endpoint = %v, want [kiwivm/getNotificationPreferences]", seen) + } + if resp.NotificationEmail != "user@example.com" { + t.Errorf("NotificationEmail = %q, want user@example.com", resp.NotificationEmail) + } + pref := resp.EmailPreferences["service"]["maintenance"] + if pref.FriendlyDescription != "Maintenance notices" { + t.Errorf("FriendlyDescription = %q, want Maintenance notices", pref.FriendlyDescription) + } + if pref.IsEnabled != 1 { + t.Errorf("IsEnabled = %d, want 1", pref.IsEnabled) + } + if pref.ChangedTimestamp != 1710000000 { + t.Errorf("ChangedTimestamp = %d, want 1710000000", pref.ChangedTimestamp) + } + if pref.SValue != "daily" { + t.Errorf("SValue = %q, want daily", pref.SValue) + } +} + +func TestClient_NewReadOnlyMethods_BWHError(t *testing.T) { + tests := []struct { + name string + endpoint string + call func(context.Context, *Client) error + }{ + { + name: "suspension details", + endpoint: "getSuspensionDetails", + call: func(ctx context.Context, c *Client) error { + _, err := c.GetSuspensionDetails(ctx) + return err + }, + }, + { + name: "policy violations", + endpoint: "getPolicyViolations", + call: func(ctx context.Context, c *Client) error { + _, err := c.GetPolicyViolations(ctx) + return err + }, + }, + { + name: "notification preferences", + endpoint: "kiwivm/getNotificationPreferences", + call: func(ctx context.Context, c *Client) error { + _, err := c.GetNotificationPreferences(ctx) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + seen := []string{} + server := newReadOnlyMockServer(t, map[string]string{ + tt.endpoint: `{"error":700005,"message":"Authentication failure"}`, + }, &seen) + defer server.Close() + + c := NewClient("invalid_key", "123456") + c.SetBaseURL(server.URL) + + err := tt.call(context.Background(), c) + if err == nil { + t.Fatal("expected BWH error") + } + if len(seen) != 1 || seen[0] != tt.endpoint { + t.Fatalf("endpoint = %v, want [%s]", seen, tt.endpoint) + } + if !IsBWHError(err) { + t.Fatalf("expected BWHError, got %T: %v", err, err) + } + bwhErr, ok := GetBWHError(err) + if !ok { + t.Fatal("failed to extract BWHError") + } + if bwhErr.Code != 700005 { + t.Errorf("BWHError.Code = %d, want 700005", bwhErr.Code) + } + }) + } +} diff --git a/pkg/client/types.go b/pkg/client/types.go index 949ffa5..44ef828 100644 --- a/pkg/client/types.go +++ b/pkg/client/types.go @@ -44,11 +44,11 @@ type ServiceInfo struct { DataNextReset int64 `json:"data_next_reset"` // Date and time of transfer counter reset (UNIX timestamp) // Network Configuration - IPAddresses []string `json:"ip_addresses"` // IPv4 addresses and IPv6 /64 subnets assigned to VPS - IPv6SitTunnelEndpoint string `json:"ipv6_sit_tunnel_endpoint"` // IPv6 SIT tunnel endpoint - PrivateIPAddresses []string `json:"private_ip_addresses"` // Private IPv4 addresses assigned to VPS - IPNullroutes []string `json:"ip_nullroutes"` // Information on IP address nullrouting during (D)DoS attacks - PlanMaxIPv6s int `json:"plan_max_ipv6s"` // Maximum number of IPv6 /64 subnets allowed by plan + IPAddresses []string `json:"ip_addresses"` // IPv4 addresses and IPv6 /64 subnets assigned to VPS + IPv6SitTunnelEndpoint string `json:"ipv6_sit_tunnel_endpoint"` // IPv6 SIT tunnel endpoint + PrivateIPAddresses []string `json:"private_ip_addresses"` // Private IPv4 addresses assigned to VPS + IPNullroutes IPNullroutes `json:"ip_nullroutes"` // Information on IP address nullrouting during (D)DoS attacks + PlanMaxIPv6s int `json:"plan_max_ipv6s"` // Maximum number of IPv6 /64 subnets allowed by plan // ISO Images ISO1 string `json:"iso1"` // Mounted image #1 @@ -76,6 +76,42 @@ type CreateSnapshotResponse struct { NotificationEmail string `json:"notificationEmail"` } +// IPNullrouteInfo contains details for a null-routed IP address. +type IPNullrouteInfo struct { + NullrouteTimestamp int64 `json:"nullroute_timestamp"` + NullrouteDurationSeconds int `json:"nullroute_duration_s"` + Log string `json:"log"` +} + +// IPNullroutes maps IP addresses to their nullroute details. +type IPNullroutes map[string]IPNullrouteInfo + +// UnmarshalJSON supports KiwiVM's empty-array response and non-empty object response. +func (n *IPNullroutes) UnmarshalJSON(data []byte) error { + if strings.TrimSpace(string(data)) == "null" { + *n = nil + return nil + } + + var routes map[string]IPNullrouteInfo + if err := json.Unmarshal(data, &routes); err == nil { + *n = routes + return nil + } + + var ips []string + if err := json.Unmarshal(data, &ips); err != nil { + return err + } + + routes = make(map[string]IPNullrouteInfo, len(ips)) + for _, ip := range ips { + routes[ip] = IPNullrouteInfo{} + } + *n = routes + return nil +} + // FlexibleInt is a type that can unmarshal both string and int from JSON type FlexibleInt struct { Value int64 @@ -220,6 +256,59 @@ type RateLimitStatus struct { RemainingPoints24H int `json:"remaining_points_24h"` // API calls remaining in 24-hour window } +// SuspensionRecord describes an outstanding service suspension issue. +type SuspensionRecord struct { + RecordID int `json:"record_id"` + Flag string `json:"flag"` + IsSoft int `json:"is_soft"` + EvidenceRecordID int `json:"evidence_record_id"` + AbusePoints int `json:"abuse_points"` +} + +// SuspensionDetailsResponse represents the response from getSuspensionDetails. +type SuspensionDetailsResponse struct { + BaseResponse + SuspensionCount int `json:"suspension_count"` + TotalAbusePoints int `json:"total_abuse_points"` + MaxAbusePoints int `json:"max_abuse_points"` + Suspensions []SuspensionRecord `json:"suspensions,omitempty"` + Evidence map[string]string `json:"evidence,omitempty"` +} + +// PolicyViolationRecord describes an unresolved policy violation. +type PolicyViolationRecord struct { + RecordID int `json:"record_id"` + Timestamp int64 `json:"timestamp"` + SuspendAt int64 `json:"suspend_at"` + Flag string `json:"flag"` + IsSoft int `json:"is_soft"` + AbusePoints int `json:"abuse_points"` + EvidenceData string `json:"evidence_data"` +} + +// PolicyViolationsResponse represents the response from getPolicyViolations. +type PolicyViolationsResponse struct { + BaseResponse + TotalAbusePoints int `json:"total_abuse_points"` + MaxAbusePoints int `json:"max_abuse_points"` + PolicyViolations []PolicyViolationRecord `json:"policy_violations,omitempty"` +} + +// NotificationPreference describes a KiwiVM e-mail notification preference. +type NotificationPreference struct { + FriendlyDescription string `json:"friendly_description"` + IsEnabled int `json:"is_enabled"` + ChangedTimestamp int64 `json:"changed_timestamp"` + SValue string `json:"s_value"` +} + +// NotificationPreferencesResponse represents notification settings and their state. +type NotificationPreferencesResponse struct { + BaseResponse + EmailPreferences map[string]map[string]NotificationPreference `json:"email_preferences"` + NotificationEmail string `json:"notificationEmail"` +} + // SshKeysResponse represents the response from getSshKeys API call type SshKeysResponse struct { BaseResponse diff --git a/pkg/client/types_test.go b/pkg/client/types_test.go index 86ab1ef..1f4754f 100644 --- a/pkg/client/types_test.go +++ b/pkg/client/types_test.go @@ -126,6 +126,67 @@ func TestBaseResponse(t *testing.T) { } } +func TestIPNullroutes_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + want IPNullroutes + }{ + { + name: "empty array", + input: `[]`, + want: IPNullroutes{}, + }, + { + name: "object keyed by IP", + input: `{ + "192.0.2.10": { + "nullroute_timestamp": 1710000000, + "nullroute_duration_s": 3600, + "log": "Nullroute active" + } + }`, + want: IPNullroutes{ + "192.0.2.10": { + NullrouteTimestamp: 1710000000, + NullrouteDurationSeconds: 3600, + Log: "Nullroute active", + }, + }, + }, + { + name: "legacy string array", + input: `["192.0.2.20"]`, + want: IPNullroutes{ + "192.0.2.20": {}, + }, + }, + { + name: "null", + input: `null`, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got IPNullroutes + if err := json.Unmarshal([]byte(tt.input), &got); err != nil { + t.Fatalf("IPNullroutes unmarshal error = %v", err) + } + + if len(got) != len(tt.want) { + t.Fatalf("IPNullroutes length = %d, want %d", len(got), len(tt.want)) + } + for ip, want := range tt.want { + if got[ip] != want { + t.Errorf("IPNullroutes[%s] = %+v, want %+v", ip, got[ip], want) + } + } + }) + } +} + func TestServiceInfo_JSONParsing(t *testing.T) { // Test parsing a minimal ServiceInfo response input := `{