Merge pull request #77 from m-mizutani/enhance/playbook-format

Refactor play scenario format

docs/test.md

@@ -91,31 +91,27 @@ test_ignore_type {
 ```
 ## Testing Action Policy
 
-Action Policy is a policy that controls the behavior of actions. As such, testing its behavior requires interactions with external services. However, using responses from external services directly in tests can be inconvenient due to constraints such as inconsistent responses or difficulty in preparing expected answers. To address this, AlertChain has implemented a "play" mode. In play mode, you can pre-define a playbook, which describes scenarios specifying how actions should respond.
+Action Policy is a policy that controls the behavior of actions. As such, testing its behavior requires interactions with external services. However, using responses from external services directly in tests can be inconvenient due to constraints such as inconsistent responses or difficulty in preparing expected answers. To address this, AlertChain has implemented a "play" mode. In play mode, you can pre-define a **Scenario**, which describes workflow specifying how actions should respond.
 
 The play mode itself is not for verifying the behavior of the policy; it only logs the execution results. However, by testing these logs using OPA/Rego, you can verify how the Action Policy behaved based on the responses obtained from each action. This achieves the "Automatic test for orchestration and automated response," which is one of the challenges in SOAR implementation.
 
 ### Playbook
 
-Here is an example of a Playbook jsonnet file:
+Here is an example of a Scenario jsonnet file:
 
 ```jsonnet
 {
-  scenarios: [
+  id: 'scenario1',
+  title: 'Test 1',
+  events: [
     {
-      id: 'scenario1',
-      title: 'Test 1',
-      events: [
-        {
-          input: import 'event/guardduty.json',
-          schema: 'aws_guardduty',
-          actions: {
-            'chatgpt.comment_alert': [
-              import 'results/chatgpt.json',
-            ],
-          },
-        },
-      ],
+      input: import 'event/guardduty.json',
+      schema: 'aws_guardduty',
+      actions: {
+        'chatgpt.comment_alert': [
+          import 'results/chatgpt.json',
+        ],
+      },
     },
   ],
   env: {
@@ -127,9 +123,11 @@ Here is an example of a Playbook jsonnet file:
 
 A scenario is composed of the following fields:
 
-- `scenarios`
-  - `id`: Specify any string, ensuring it is unique within the playbook. This serves as a key to identify the scenario when writing tests using Rego.
-  - `event`: This field is used to import the alert data required for the scenario.
+
+- `id`: Specify any string, ensuring it is unique within the playbook. This serves as a key to identify the scenario when writing tests using Rego.
+- `title`: Specify any string. This is used to describe the scenario for human readability.
+- `events`: This describes scenarios for each event.
+  - `input`: This field specifies the event data to be used for the scenario.
   - `schema`: This field specifies the schema to be used for the scenario.
   - `actions`: This field contains the expected results for each action involved in the scenario. The results are defined as key-value pairs, where the key represents the action Name and the value is an array of expected responses for that action.
 - `env`: Environment variables that will be used in play mode.

examples/test/policy/action.rego

@@ -4,7 +4,7 @@ run[res] {
 	input.alert.source == "aws"
 	res := {
 		"id": "ask-gpt",
-		"uses": "chatgpt.comment_alert",
+		"uses": "chatgpt.query",
 		"args": {"secret_api_key": input.env.CHATGPT_API_KEY},
 	}
 }

examples/test/playbook.jsonnet → examples/test/scenarios/scenario.jsonnet

@@ -1,17 +1,18 @@
 {
-  scenarios: [
+  id: 'scenario1',
+  title: 'Test 1',
+  events: [
     {
-      id: 'scenario1',
-      title: 'Test 1',
-      event: import 'event/guardduty.json',
+      input: import '../event/guardduty.json',
       schema: 'aws_guardduty',
-      results: {
-        'chatgpt.comment_alert': [
-          import 'results/chatgpt.json',
+      actions: {
+        'chatgpt.query': [
+            import '../results/chatgpt.json',
         ],
       },
     },
   ],
+
   env: {
     CHATGPT_API_KEY: 'test_api_key_xxxxxxxxxx',
     SLACK_WEBHOOK_URL: 'https://hooks.slack.com/services/xxxxxxxxx',

go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/m-mizutani/clog v0.0.4
 	github.com/m-mizutani/goerr v0.1.12-0.20240113043120-3feee990f78c
 	github.com/m-mizutani/gots v0.0.0-20230529013424-0639119b2cdd
-	github.com/m-mizutani/gt v0.0.6-0.20230708234934-97ecdb8cc874
+	github.com/m-mizutani/gt v0.0.9-0.20240121024259-2c2e7bf7b4f8
 	github.com/m-mizutani/masq v0.1.5
 	github.com/open-policy-agent/opa v0.60.0
 	github.com/opsgenie/opsgenie-go-sdk-v2 v1.2.22

go.sum

@@ -205,8 +205,8 @@ github.com/m-mizutani/goerr v0.1.12-0.20240113043120-3feee990f78c h1:HxLB9f6oLVC
 github.com/m-mizutani/goerr v0.1.12-0.20240113043120-3feee990f78c/go.mod h1:64HHjaK/ZjCy3VMaqrcZvinirVZkIBUxU21ml3WgMU4=
 github.com/m-mizutani/gots v0.0.0-20230529013424-0639119b2cdd h1:mmuUG300WqGOvGNIIK3ELod2LpZZBV94fEgXK36xM0o=
 github.com/m-mizutani/gots v0.0.0-20230529013424-0639119b2cdd/go.mod h1:MDYrqsaKL6z6djDZpy6admhX6GOOCvyST+c0VnLbT4w=
-github.com/m-mizutani/gt v0.0.6-0.20230708234934-97ecdb8cc874 h1:zGmsSQCkuHoL9SUcDJT204sYD+ugL2gWgqXRZy7ZFO4=
-github.com/m-mizutani/gt v0.0.6-0.20230708234934-97ecdb8cc874/go.mod h1:0MPYSfGBLmYjTduzADVmIqD58ELQ5IfBFiK/f0FmB3k=
+github.com/m-mizutani/gt v0.0.9-0.20240121024259-2c2e7bf7b4f8 h1:dt+Xn8q0rrRmqAuwwp4gNDnfO5cTW7uROnLx9eLZg7o=
+github.com/m-mizutani/gt v0.0.9-0.20240121024259-2c2e7bf7b4f8/go.mod h1:0MPYSfGBLmYjTduzADVmIqD58ELQ5IfBFiK/f0FmB3k=
 github.com/m-mizutani/masq v0.1.5 h1:+ebFJ9gdIZdiNA0X3Cn8MAXgfH4Z6m0cLJE9UNcKQKk=
 github.com/m-mizutani/masq v0.1.5/go.mod h1:42/bKhlCNIQjmh3KBypeuh6iOvxNfUIlrZD1i0amEoc=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=

main_test.go

@@ -3,15 +3,18 @@ package main_test
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"io"
 	"net/http"
 	"testing"
 	"time"
 
 	"github.com/m-mizutani/alertchain/pkg/controller/cli"
+	"github.com/m-mizutani/alertchain/pkg/domain/model"
 	"github.com/m-mizutani/gt"
 )
 
-func TestE2E(t *testing.T) {
+func TestServe(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -63,3 +66,43 @@ func TestE2E(t *testing.T) {
 	send(t) // 4
 	gt.N(t, called).Equal(2)
 }
+
+func TestPlay(t *testing.T) {
+	ctx := context.Background()
+	args := []string{
+		"alertchain",
+		"-l", "debug",
+		"play",
+		"-d", "examples/test/policy",
+		"-s", "examples/test/scenarios",
+		"-o", "examples/test/output",
+	}
+	gt.NoError(t, cli.New().Run(ctx, args))
+
+	gt.F(t, "examples/test/output/scenario1/data.json").Reader(func(t testing.TB, r io.Reader) {
+		var data model.ScenarioLog
+		gt.NoError(t, json.NewDecoder(r).Decode(&data))
+		gt.Equal(t, data.ID, "scenario1")
+		gt.Equal(t, data.Title, "Test 1")
+		gt.A(t, data.Results).Length(1).
+			At(0, func(t testing.TB, v *model.PlayLog) {
+				gt.Equal(t, v.Alert.Title, "Trojan:EC2/DropPoint!DNS")
+
+				gt.A(t, v.Actions).Length(2).
+					At(0, func(t testing.TB, v *model.ActionLog) {
+						gt.Equal(t, v.Seq, 0)
+						gt.A(t, v.Run).Length(1).
+							At(0, func(t testing.TB, v model.Action) {
+								gt.Equal(t, v.Uses, "chatgpt.query")
+							})
+					}).
+					At(1, func(t testing.TB, v *model.ActionLog) {
+						gt.Equal(t, v.Seq, 1)
+						gt.A(t, v.Run).Length(1).
+							At(0, func(t testing.TB, v model.Action) {
+								gt.Equal(t, v.Uses, "slack.post")
+							})
+					})
+			})
+	})
+}

pkg/chain/testdata/play_workflow/playbook.jsonnet

@@ -9,7 +9,7 @@
             class: 'threat',
           },
           schema: 'my_test',
-          results: {
+          actions: {
             mock: [
               {
                 index: 'first',

pkg/controller/cli/play.go

@@ -18,7 +18,7 @@ import (
 
 func cmdPlay() *cli.Command {
 	var (
-		playbookPath string
+		scenarioPath string
 		outDir       string
 		scenarioIDs  cli.StringSlice
 
@@ -27,12 +27,12 @@ func cmdPlay() *cli.Command {
 
 	flags := []cli.Flag{
 		&cli.StringFlag{
-			Name:        "playbook",
-			Aliases:     []string{"b"},
-			Usage:       "playbook file",
-			EnvVars:     []string{"ALERTCHAIN_PLAYBOOK"},
+			Name:        "scenario",
+			Aliases:     []string{"s"},
+			Usage:       "scenario directory",
+			EnvVars:     []string{"ALERTCHAIN_SCENARIO"},
 			Required:    true,
-			Destination: &playbookPath,
+			Destination: &scenarioPath,
 		},
 		&cli.StringFlag{
 			Name:        "output",
@@ -43,10 +43,10 @@ func cmdPlay() *cli.Command {
 			Value:       "./output",
 		},
 		&cli.StringSliceFlag{
-			Name:        "scenario",
-			Aliases:     []string{"s"},
-			Usage:       "scenario ID to play. If not specified, all scenarios are played",
-			EnvVars:     []string{"ALERTCHAIN_SCENARIO"},
+			Name:        "target",
+			Aliases:     []string{"t"},
+			Usage:       "Target scenario ID to play. If not specified, all scenarios are played",
+			EnvVars:     []string{"ALERTCHAIN_TARGET"},
 			Destination: &scenarioIDs,
 		},
 	}
@@ -60,16 +60,41 @@ func cmdPlay() *cli.Command {
 
 		Action: func(c *cli.Context) error {
 			// Load playbook
-			var playbook model.Playbook
-			if err := model.ParsePlaybook(playbookPath, os.ReadFile, &playbook); err != nil {
-				return goerr.Wrap(err, "failed to parse playbook")
+			scenarioFiles := make([]string, 0)
+			err := filepath.Walk(scenarioPath, func(path string, info os.FileInfo, err error) error {
+				if err != nil {
+					return err
+				}
+				if !info.IsDir() && filepath.Ext(path) == ".jsonnet" {
+					scenarioFiles = append(scenarioFiles, path)
+				}
+				return nil
+			})
+			if err != nil {
+				return goerr.Wrap(err, "failed to walk through playbook directory")
 			}
 
 			ctx := model.NewContext(
 				model.WithBase(c.Context),
 				model.WithCLI(),
 			)
-			ctx.Logger().Info("starting alertchain with play mode", slog.Any("playbook", playbookPath))
+
+			var playbook model.Playbook
+			for _, scenarioFile := range scenarioFiles {
+				ctx.Logger().Debug("Load scenario", slog.String("file", scenarioFile))
+				s, err := model.ParseScenario(scenarioFile, os.ReadFile)
+				if err != nil {
+					return goerr.Wrap(err, "failed to parse playbook")
+				}
+
+				playbook.Scenarios = append(playbook.Scenarios, s)
+			}
+
+			if err := playbook.Validate(); err != nil {
+				return err
+			}
+
+			ctx.Logger().Info("starting alertchain with play mode", slog.Any("scenario dir", scenarioPath))
 
 			targets := make(map[types.ScenarioID]struct{})
 			for _, id := range scenarioIDs.Value() {
@@ -81,7 +106,7 @@ func cmdPlay() *cli.Command {
 					continue
 				}
 
-				if err := playScenario(ctx, s, &policyCfg, outDir, playbook.Env); err != nil {
+				if err := playScenario(ctx, s, &policyCfg, outDir); err != nil {
 					return err
 				}
 			}
@@ -99,7 +124,7 @@ func (x *actionMockWrapper) GetResult(name types.ActionName) any {
 	return x.ev.GetResult(name)
 }
 
-func playScenario(ctx *model.Context, scenario *model.Scenario, cfg *config.Policy, outDir string, envVars types.EnvVars) error {
+func playScenario(ctx *model.Context, scenario *model.Scenario, cfg *config.Policy, outDir string) error {
 	ctx.Logger().Debug("Start scenario", slog.Any("scenario", scenario))
 
 	w, err := openLogFile(outDir, string(scenario.ID))
@@ -119,9 +144,9 @@ func playScenario(ctx *model.Context, scenario *model.Scenario, cfg *config.Poli
 		core.WithActionMock(mockWrapper),
 	}
 
-	if envVars != nil {
+	if scenario.Env != nil {
 		options = append(options, core.WithEnv(func() types.EnvVars {
-			return envVars
+			return scenario.Env
 		}))
 	}
 

pkg/domain/model/alert.go

@@ -34,11 +34,11 @@ type Alert struct {
 	AlertMetaData
 	ID        types.AlertID `json:"id"`
 	Schema    types.Schema  `json:"schema"`
-	Data      any           `json:"data"`
+	Data      any           `json:"data,omitempty"`
 	CreatedAt time.Time     `json:"created_at"`
 
 	// Raw is a JSON string of Data. The field will be redacted by masq because of verbosity.
-	Raw string `json:"raw" masq:"quiet"`
+	Raw string `json:"raw,omitempty" masq:"quiet"`
 }
 
 func (x Alert) Copy() Alert {

pkg/domain/model/playbook.go

@@ -10,8 +10,7 @@ import (
 )
 
 type Playbook struct {
-	Scenarios []*Scenario   `json:"scenarios"`
-	Env       types.EnvVars `json:"env"`
+	Scenarios []*Scenario `json:"scenarios"`
 }
 
 func (x *Playbook) Validate() error {
@@ -40,7 +39,7 @@ func (x *Playbook) Validate() error {
 type Event struct {
 	Input   any                        `json:"input"`
 	Schema  types.Schema               `json:"schema"`
-	Results map[types.ActionName][]any `json:"results"`
+	Actions map[types.ActionName][]any `json:"actions"`
 
 	actionIndex map[types.ActionName]int
 }
@@ -49,6 +48,7 @@ type Scenario struct {
 	ID     types.ScenarioID    `json:"id"`
 	Title  types.ScenarioTitle `json:"title"`
 	Events []Event             `json:"events"`
+	Env    types.EnvVars       `json:"env"`
 }
 
 func (x *Scenario) Validate() error {
@@ -89,12 +89,12 @@ func (x *Event) GetResult(actionName types.ActionName) any {
 	if !ok {
 		idx = 0
 	}
-	if len(x.Results[actionName]) <= idx {
+	if len(x.Actions[actionName]) <= idx {
 		return nil
 	}
 	x.actionIndex[actionName] = idx + 1
 
-	return x.Results[actionName][idx]
+	return x.Actions[actionName][idx]
 }
 
 type embedImporter struct {
@@ -129,3 +129,24 @@ func ParsePlaybook(entryFile string, readFile ReadFile, book *Playbook) error {
 
 	return book.Validate()
 }
+
+func ParseScenario(entryFile string, readFile ReadFile) (*Scenario, error) {
+	vm := jsonnet.MakeVM()
+	vm.Importer(&embedImporter{readFile: readFile})
+
+	raw, err := vm.EvaluateFile(entryFile)
+	if err != nil {
+		return nil, goerr.Wrap(err, "evaluating scenario jsonnet")
+	}
+
+	var scenario Scenario
+	if err := json.Unmarshal([]byte(raw), &scenario); err != nil {
+		return nil, goerr.Wrap(err, "unmarshal scenario by jsonnet")
+	}
+
+	if err := scenario.Validate(); err != nil {
+		return nil, err
+	}
+
+	return &scenario, nil
+}