diff --git a/tavern/internal/mcp/export_test.go b/tavern/internal/mcp/export_test.go new file mode 100644 index 000000000..4d8a1232a --- /dev/null +++ b/tavern/internal/mcp/export_test.go @@ -0,0 +1,12 @@ +package mcp + +// Export unexported types and functions for testing + +type ContextKey = contextKey + +var ( + HandleListHosts = handleListHosts + HandleListQuests = handleListQuests + HandleListTomes = handleListTomes + HandleQuestOutput = handleQuestOutput +) diff --git a/tavern/internal/mcp/tool_list_hosts_test.go b/tavern/internal/mcp/tool_list_hosts_test.go new file mode 100644 index 000000000..34904fa40 --- /dev/null +++ b/tavern/internal/mcp/tool_list_hosts_test.go @@ -0,0 +1,116 @@ +package mcp_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "realm.pub/tavern/internal/c2/c2pb" + "realm.pub/tavern/internal/ent/tag" + tavernmcp "realm.pub/tavern/internal/mcp" +) + +func TestHandleListHosts(t *testing.T) { + client := setupTestDB(t) + defer client.Close() + + ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client) + + // Create tags + tag1 := client.Tag.Create().SetName("web").SetKind(tag.KindService).SaveX(ctx) + tag2 := client.Tag.Create().SetName("db").SetKind(tag.KindService).SaveX(ctx) + + // Create mock hosts + host1 := client.Host.Create(). + SetIdentifier("host-1"). + SetPlatform(c2pb.Host_PLATFORM_LINUX). + SetName("web-server"). + SetPrimaryIP("192.168.1.10"). + SetExternalIP("203.0.113.10"). + SetLastSeenAt(time.Now()). + AddTags(tag1). + SaveX(ctx) + + // Create mock beacons and attach them to host1 + client.Beacon.Create(). + SetHost(host1). + SetName("beacon-1"). + SetTransport(c2pb.Transport_TRANSPORT_HTTP1). + SetPrincipal("admin"). + SaveX(ctx) + + host2 := client.Host.Create(). + SetIdentifier("host-2"). + SetPlatform(c2pb.Host_PLATFORM_WINDOWS). + SetName("db-server"). + SetPrimaryIP("192.168.1.20"). + SetExternalIP("203.0.113.20"). + SetLastSeenAt(time.Now()). + AddTags(tag2). + SaveX(ctx) + _ = host2 + + req := mcp.CallToolRequest{} + res, err := tavernmcp.HandleListHosts(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + require.False(t, res.IsError) + + var results []map[string]interface{} + textContent := res.Content[0].(mcp.TextContent).Text + err = json.Unmarshal([]byte(textContent), &results) + require.NoError(t, err) + + assert.Len(t, results, 2) + + // Verify host 1 + var h1Res map[string]interface{} + var h2Res map[string]interface{} + if results[0]["id"].(float64) == float64(host1.ID) { + h1Res = results[0] + h2Res = results[1] + } else { + h1Res = results[1] + h2Res = results[0] + } + + assert.Equal(t, "host-1", h1Res["identifier"]) + assert.Equal(t, "web-server", h1Res["name"]) + assert.Equal(t, "PLATFORM_LINUX", h1Res["platform"]) + assert.Equal(t, "192.168.1.10", h1Res["primaryIP"]) + assert.Equal(t, "203.0.113.10", h1Res["externalIP"]) + + tags1 := h1Res["tags"].([]interface{}) + require.Len(t, tags1, 1) + tagRes := tags1[0].(map[string]interface{}) + assert.Equal(t, "web", tagRes["name"]) + assert.Equal(t, "service", tagRes["kind"]) + + beacons1 := h1Res["beacons"].([]interface{}) + require.Len(t, beacons1, 1) + beaconRes := beacons1[0].(map[string]interface{}) + assert.Equal(t, "beacon-1", beaconRes["name"]) + assert.Equal(t, "admin", beaconRes["principal"]) + + // Verify host 2 + assert.Equal(t, "host-2", h2Res["identifier"]) + assert.Equal(t, "db-server", h2Res["name"]) + assert.Equal(t, "PLATFORM_WINDOWS", h2Res["platform"]) + assert.Equal(t, "192.168.1.20", h2Res["primaryIP"]) + assert.Equal(t, "203.0.113.20", h2Res["externalIP"]) + tags2 := h2Res["tags"].([]interface{}) + require.Len(t, tags2, 1) + tagRes2 := tags2[0].(map[string]interface{}) + assert.Equal(t, "db", tagRes2["name"]) + assert.Equal(t, "service", tagRes2["kind"]) + + // test no client error + res, err = tavernmcp.HandleListHosts(context.Background(), req) + require.NoError(t, err) + require.True(t, res.IsError) + assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client") +} diff --git a/tavern/internal/mcp/tool_list_quests_test.go b/tavern/internal/mcp/tool_list_quests_test.go new file mode 100644 index 000000000..38abc7087 --- /dev/null +++ b/tavern/internal/mcp/tool_list_quests_test.go @@ -0,0 +1,105 @@ +package mcp_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "realm.pub/tavern/internal/ent/tome" + tavernmcp "realm.pub/tavern/internal/mcp" +) + +func TestHandleListQuests(t *testing.T) { + client := setupTestDB(t) + defer client.Close() + + ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client) + + // Create user + user1 := client.User.Create(). + SetOauthID("oauth-1"). + SetName("Alice"). + SetPhotoURL("http://example.com/photo"). + SaveX(ctx) + + // Create mock tomes + tome1 := client.Tome.Create(). + SetName("discovery-tome"). + SetTactic(tome.TacticDISCOVERY). + SetDescription("Discover things"). + SetParamDefs(`[{"name": "p1", "type":"string"}]`). + SetAuthor("Alice"). + SetEldritch("console.log('hello')"). + SaveX(ctx) + + tome2 := client.Tome.Create(). + SetName("exec-tome"). + SetTactic(tome.TacticEXECUTION). + SetDescription("Execute things"). + SetParamDefs(`[]`). + SetAuthor("Bob"). + SetEldritch("console.log('world')"). + SaveX(ctx) + + quest1 := client.Quest.Create(). + SetName("test-quest-1"). + SetParameters(`{"p1": "v1"}`). + SetTome(tome1). + SetCreator(user1). + SaveX(ctx) + + quest2 := client.Quest.Create(). + SetName("test-quest-2"). + SetParameters(`{}`). + SetTome(tome2). + SaveX(ctx) + _ = quest2 + + req := mcp.CallToolRequest{} + res, err := tavernmcp.HandleListQuests(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + require.False(t, res.IsError) + + var results []map[string]interface{} + textContent := res.Content[0].(mcp.TextContent).Text + err = json.Unmarshal([]byte(textContent), &results) + require.NoError(t, err) + + assert.Len(t, results, 2) + + var q1Res map[string]interface{} + var q2Res map[string]interface{} + if results[0]["id"].(float64) == float64(quest1.ID) { + q1Res = results[0] + q2Res = results[1] + } else { + q1Res = results[1] + q2Res = results[0] + } + + assert.Equal(t, "test-quest-1", q1Res["name"]) + assert.Equal(t, "Alice", q1Res["creator"]) + assert.Equal(t, `{"p1": "v1"}`, q1Res["parameters"]) + assert.Equal(t, "discovery-tome", q1Res["tomeName"]) + assert.Equal(t, "DISCOVERY", q1Res["tomeTactic"]) + assert.Equal(t, "Discover things", q1Res["tomeDescription"]) + assert.NotEmpty(t, q1Res["createdAt"]) + + assert.Equal(t, "test-quest-2", q2Res["name"]) + assert.NotContains(t, q2Res, "creator") + assert.Equal(t, `{}`, q2Res["parameters"]) + assert.Equal(t, "exec-tome", q2Res["tomeName"]) + assert.Equal(t, "EXECUTION", q2Res["tomeTactic"]) + assert.Equal(t, "Execute things", q2Res["tomeDescription"]) + assert.NotEmpty(t, q2Res["createdAt"]) + + // test no client error + res, err = tavernmcp.HandleListQuests(context.Background(), req) + require.NoError(t, err) + require.True(t, res.IsError) + assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client") +} diff --git a/tavern/internal/mcp/tool_list_tomes_test.go b/tavern/internal/mcp/tool_list_tomes_test.go new file mode 100644 index 000000000..41499d8ea --- /dev/null +++ b/tavern/internal/mcp/tool_list_tomes_test.go @@ -0,0 +1,79 @@ +package mcp_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "realm.pub/tavern/internal/ent/tome" + tavernmcp "realm.pub/tavern/internal/mcp" +) + +func TestHandleListTomes(t *testing.T) { + client := setupTestDB(t) + defer client.Close() + + ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client) + + // Create mock tomes + tome1 := client.Tome.Create(). + SetName("discovery-tome"). + SetTactic(tome.TacticDISCOVERY). + SetDescription("Discover things"). + SetParamDefs(`[{"name": "p1", "type":"string"}]`). + SetAuthor("Alice"). + SetEldritch("console.log('hello')"). + SaveX(ctx) + + tome2 := client.Tome.Create(). + SetName("exec-tome"). + SetTactic(tome.TacticEXECUTION). + SetDescription("Execute things"). + SetParamDefs(`[]`). + SetAuthor("Bob"). + SetEldritch("console.log('world')"). + SaveX(ctx) + _ = tome2 + + req := mcp.CallToolRequest{} + res, err := tavernmcp.HandleListTomes(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + require.False(t, res.IsError) + + var results []map[string]interface{} + textContent := res.Content[0].(mcp.TextContent).Text + err = json.Unmarshal([]byte(textContent), &results) + require.NoError(t, err) + + assert.Len(t, results, 2) + + var t1Res map[string]interface{} + var t2Res map[string]interface{} + if results[0]["id"].(float64) == float64(tome1.ID) { + t1Res = results[0] + t2Res = results[1] + } else { + t1Res = results[1] + t2Res = results[0] + } + + assert.Equal(t, "discovery-tome", t1Res["name"]) + assert.Equal(t, "DISCOVERY", t1Res["tactic"]) + assert.Equal(t, "Discover things", t1Res["description"]) + assert.Equal(t, `[{"name": "p1", "type":"string"}]`, t1Res["paramDefs"]) + + assert.Equal(t, "exec-tome", t2Res["name"]) + assert.Equal(t, "EXECUTION", t2Res["tactic"]) + assert.Equal(t, "Execute things", t2Res["description"]) + assert.Equal(t, `[]`, t2Res["paramDefs"]) + + // test no client error + res, err = tavernmcp.HandleListTomes(context.Background(), req) + require.NoError(t, err) + require.True(t, res.IsError) + assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client") +} diff --git a/tavern/internal/mcp/tool_quest_output_test.go b/tavern/internal/mcp/tool_quest_output_test.go new file mode 100644 index 000000000..e67475a1e --- /dev/null +++ b/tavern/internal/mcp/tool_quest_output_test.go @@ -0,0 +1,118 @@ +package mcp_test + +import ( + "context" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "realm.pub/tavern/internal/c2/c2pb" + "realm.pub/tavern/internal/ent/tome" + tavernmcp "realm.pub/tavern/internal/mcp" +) + +func TestHandleQuestOutput(t *testing.T) { + client := setupTestDB(t) + defer client.Close() + + ctx := context.WithValue(context.Background(), tavernmcp.ContextKey{}, client) + + host1 := client.Host.Create(). + SetIdentifier("host-1"). + SetPlatform(c2pb.Host_PLATFORM_LINUX). + SetName("web-server"). + SetPrimaryIP("192.168.1.10"). + SetExternalIP("203.0.113.10"). + SetLastSeenAt(time.Now()). + SaveX(ctx) + + beacon1 := client.Beacon.Create(). + SetHost(host1). + SetName("beacon-1"). + SetTransport(c2pb.Transport_TRANSPORT_HTTP1). + SetPrincipal("admin"). + SaveX(ctx) + + tome1 := client.Tome.Create(). + SetName("discovery-tome"). + SetTactic(tome.TacticDISCOVERY). + SetDescription("Discover things"). + SetParamDefs(`[]`). + SetAuthor("Alice"). + SetEldritch("console.log('hello')"). + SaveX(ctx) + + quest1 := client.Quest.Create(). + SetName("test-quest-1"). + SetParameters(`{}`). + SetTome(tome1). + SaveX(ctx) + + client.Task.Create(). + SetQuest(quest1). + SetBeacon(beacon1). + SetExecFinishedAt(time.Now()). + SetOutput("Output from task 1"). + SaveX(ctx) + + client.Task.Create(). + SetQuest(quest1). + SetBeacon(beacon1). + SetExecFinishedAt(time.Now()). + SetOutput("Output from task 2"). + SaveX(ctx) + + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{ + "ids": []interface{}{fmt.Sprintf("%d", quest1.ID)}, + } + + res, err := tavernmcp.HandleQuestOutput(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + require.False(t, res.IsError) + + var results []map[string]interface{} + textContent := res.Content[0].(mcp.TextContent).Text + err = json.Unmarshal([]byte(textContent), &results) + require.NoError(t, err) + + assert.Len(t, results, 1) + + q1Res := results[0] + assert.Equal(t, float64(quest1.ID), q1Res["id"]) + assert.Equal(t, "test-quest-1", q1Res["name"]) + + tasks := q1Res["tasks"].([]interface{}) + assert.Len(t, tasks, 2) + + task1 := tasks[0].(map[string]interface{}) + assert.Equal(t, "Output from task 1", task1["output"]) + beaconInfo := task1["beacon"].(map[string]interface{}) + assert.Equal(t, float64(beacon1.ID), beaconInfo["id"]) + assert.Equal(t, "beacon-1", beaconInfo["name"]) + assert.Equal(t, float64(host1.ID), beaconInfo["hostId"]) + assert.Equal(t, "web-server", beaconInfo["hostName"]) + + task2 := tasks[1].(map[string]interface{}) + assert.Equal(t, "Output from task 2", task2["output"]) + + // Test invalid ids format + req.Params.Arguments = map[string]interface{}{ + "ids": "not-an-array", + } + res, err = tavernmcp.HandleQuestOutput(ctx, req) + require.NoError(t, err) + require.True(t, res.IsError) + assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "invalid ids") + + // test no client error + res, err = tavernmcp.HandleQuestOutput(context.Background(), req) + require.NoError(t, err) + require.True(t, res.IsError) + assert.Contains(t, res.Content[0].(mcp.TextContent).Text, "internal error: no database client") +}