diff --git a/go.mod b/go.mod index d40af98c..415383a0 100644 --- a/go.mod +++ b/go.mod @@ -229,7 +229,7 @@ require ( golang.org/x/crypto v0.11.0 // indirect golang.org/x/mod v0.12.0 // indirect golang.org/x/net v0.12.0 // indirect - golang.org/x/sync v0.3.0 // indirect + golang.org/x/sync v0.3.0 golang.org/x/sys v0.10.0 // indirect golang.org/x/term v0.10.0 // indirect golang.org/x/text v0.11.0 // indirect diff --git a/plugins/extractors/http/execute_script.go b/plugins/extractors/http/execute_script.go index 2b91d72c..cf14c92b 100644 --- a/plugins/extractors/http/execute_script.go +++ b/plugins/extractors/http/execute_script.go @@ -2,6 +2,7 @@ package http import ( "context" + "encoding/json" "errors" "fmt" "reflect" @@ -19,8 +20,7 @@ import ( "google.golang.org/protobuf/proto" ) -func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plugins.Emit) error { - scriptCfg := e.config.Script +func (e *Extractor) executeScript(ctx context.Context, res interface{}, scriptCfg Script, emit plugins.Emit) error { s, err := tengoutil.NewSecureScript( ([]byte)(scriptCfg.Source), e.scriptGlobals(ctx, res, emit), ) @@ -40,12 +40,23 @@ func (e *Extractor) executeScript(ctx context.Context, res interface{}, emit plu return fmt.Errorf("run: %w", err) } + err = e.convertTengoObjToRequest(c.Get("request").Value()) + if err != nil { + return err + } + return nil } func (e *Extractor) scriptGlobals(ctx context.Context, res interface{}, emit plugins.Emit) map[string]interface{} { + req, err := e.convertRequestToTengoObj() + if err != nil { + e.logger.Error(err.Error()) + } + return map[string]interface{}{ "recipe_scope": &tengo.String{Value: e.UrnScope}, + "request": req, "response": res, "new_asset": &tengo.UserFunction{ Name: "new_asset", @@ -68,6 +79,30 @@ func (e *Extractor) scriptGlobals(ctx context.Context, res interface{}, emit plu } } +func (e *Extractor) convertTengoObjToRequest(obj interface{}) error { + r, err := json.Marshal(obj) + if err != nil { + return err + } + err = json.Unmarshal(r, &e.config.Request) + if err != nil { + return err + } + return nil +} +func (e *Extractor) convertRequestToTengoObj() (tengo.Object, error) { + var res map[string]interface{} + r, err := json.Marshal(e.config.Request) + if err != nil { + return nil, err + } + err = json.Unmarshal(r, &res) + if err != nil { + return nil, err + } + return tengo.FromInterface(res) +} + func newAssetWrapper() tengo.CallableFunc { typeURLs := knownTypeURLs() return func(args ...tengo.Object) (tengo.Object, error) { diff --git a/plugins/extractors/http/http_extractor.go b/plugins/extractors/http/http_extractor.go index c3ce5917..a2483566 100644 --- a/plugins/extractors/http/http_extractor.go +++ b/plugins/extractors/http/http_extractor.go @@ -29,17 +29,20 @@ func init() { //go:embed README.md var summary string +type Script struct { + Engine string `mapstructure:"engine" validate:"required,oneof=tengo"` + Source string `mapstructure:"source" validate:"required"` + MaxAllocs int64 `mapstructure:"max_allocs" validate:"gt=100" default:"5000"` + MaxConstObjects int `mapstructure:"max_const_objects" validate:"gt=10" default:"500"` +} + // Config holds the set of configuration for the HTTP extractor. type Config struct { Request RequestConfig `mapstructure:"request"` SuccessCodes []int `mapstructure:"success_codes" validate:"dive,gte=200,lt=300" default:"[200]"` Concurrency int `mapstructure:"concurrency" validate:"gte=1,lte=100" default:"5"` - Script struct { - Engine string `mapstructure:"engine" validate:"required,oneof=tengo"` - Source string `mapstructure:"source" validate:"required"` - MaxAllocs int64 `mapstructure:"max_allocs" validate:"gt=100" default:"5000"` - MaxConstObjects int `mapstructure:"max_const_objects" validate:"gt=10" default:"500"` - } `mapstructure:"script"` + Script Script `mapstructure:"script"` + BeforeScript *Script `mapstructure:"before_script"` } type RequestConfig struct { @@ -122,12 +125,17 @@ func (e *Extractor) Init(ctx context.Context, config plugins.Config) error { // executes the script. The script has access to the response and can use the // same to 'emit' assets from within the script. func (e *Extractor) Extract(ctx context.Context, emit plugins.Emit) error { + if e.config.BeforeScript != nil { + if err := e.executeScript(ctx, nil, *e.config.BeforeScript, emit); err != nil { + return fmt.Errorf("http extractor: execute script: %w", err) + } + } res, err := e.executeRequest(ctx, e.config.Request) if err != nil { return fmt.Errorf("http extractor: execute request: %w", err) } - if err := e.executeScript(ctx, res, emit); err != nil { + if err := e.executeScript(ctx, res, e.config.Script, emit); err != nil { return fmt.Errorf("http extractor: execute script: %w", err) } diff --git a/plugins/extractors/http/http_extractor_test.go b/plugins/extractors/http/http_extractor_test.go index 4f6a8509..84ea6768 100644 --- a/plugins/extractors/http/http_extractor_test.go +++ b/plugins/extractors/http/http_extractor_test.go @@ -202,6 +202,66 @@ func TestExtract(t *testing.T) { testutils.Respond(t, w, http.StatusOK, `[]`) }, }, + { + name: "MatchRequestBeforeScript", + rawCfg: map[string]interface{}{ + "before_script": map[string]interface{}{ + "engine": "tengo", + "source": heredoc.Doc(` + fmt := import("fmt") + reqs := [] + reqs = append(reqs, { + url: "{{serverURL}}/token", + method: "GET", + content_type: "application/json", + accept: "application/json", + timeout: "5s" + }) + + responses := execute_request(reqs...) + for r in responses { + if is_error(r) { + continue + } + request.Headers = {"Authorization": format("Bearer %s", r.body.data.token)} + } + `), + "max_allocs": 5000, + "max_const_objects": 500, + }, + "request": map[string]interface{}{ + "url": "{{serverURL}}/api/v2/endpoint", + "content_type": "application/json", + "accept": "application/json", + }, + "script": map[string]interface{}{ + "engine": "tengo", + "source": "// do nothing", + }, + }, + handler: func(t *testing.T, w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/token": + testutils.Respond(t, w, http.StatusOK, + `{"header":{"process_time":0.130462645,"messages":[],"error_code":""},"data":{"token":"testToken123123","service_id":1,"expires_at":1711537963}}`, + ) + case "/api/v2/endpoint": + assert.Equal(t, r.Method, http.MethodGet) + assert.Equal(t, r.URL.Path, "/api/v2/endpoint") + assert.Equal(t, r.URL.RawQuery, "") + h := r.Header + assert.Equal(t, "", h.Get("Content-Type")) + assert.Equal(t, "Bearer testToken123123", h.Get("Authorization")) + assert.Equal(t, "application/json", h.Get("Accept")) + data, err := io.ReadAll(r.Body) + assert.NoError(t, err) + assert.Empty(t, data) + testutils.Respond(t, w, http.StatusOK, `[]`) + default: + t.Error("Unexpected HTTP call on", r.URL.Path) + } + }, + }, { name: "MatchRequestAdvanced", rawCfg: map[string]interface{}{ @@ -890,4 +950,8 @@ func replaceServerURL(cfg map[string]interface{}, serverURL string) { reqCfg["url"] = strings.Replace(reqCfg["url"].(string), "{{serverURL}}", serverURL, 1) scriptCfg := cfg["script"].(map[string]interface{}) scriptCfg["source"] = strings.Replace(scriptCfg["source"].(string), "{{serverURL}}", serverURL, -1) + beforeScriptCfg, ok := cfg["before_script"].(map[string]interface{}) + if ok { + beforeScriptCfg["source"] = strings.Replace(beforeScriptCfg["source"].(string), "{{serverURL}}", serverURL, -1) + } }