Skip to content

Commit

Permalink
Merge pull request #403 from projectdiscovery/add-yaml-based-workflow
Browse files Browse the repository at this point in the history
Add yaml based workflow generator
  • Loading branch information
ehsandeep committed Nov 12, 2020
2 parents 87d2c17 + 302ca4d commit d4e03a0
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 4 deletions.
2 changes: 2 additions & 0 deletions v2/internal/runner/processor.go
Original file line number Diff line number Diff line change
Expand Up @@ -292,12 +292,14 @@ func (r *Runner) preloadWorkflowTemplates(p *progress.Progress, workflow *workfl
ProxySocksURL: r.options.ProxySocksURL,
CustomHeaders: r.options.CustomHeaders,
CookieJar: jar,
TraceLog: r.traceLog,
}
} else if len(t.RequestsDNS) > 0 {
template.DNSOptions = &executer.DNSOptions{
Debug: r.options.Debug,
Template: t,
Writer: r.output,
TraceLog: r.traceLog,
}
}
if template.DNSOptions != nil || template.HTTPOptions != nil {
Expand Down
9 changes: 6 additions & 3 deletions v2/pkg/workflows/compile.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
package workflows

import (
"errors"
"os"

"github.com/pkg/errors"
"gopkg.in/yaml.v2"
)

Expand All @@ -22,11 +22,14 @@ func Parse(file string) (*Workflow, error) {
return nil, err
}

if len(workflow.Workflows) > 0 {
if err := workflow.generateLogicFromWorkflows(); err != nil {
return nil, errors.Wrap(err, "could not generate workflow")
}
}
if workflow.Logic == "" {
return nil, errors.New("no logic provided")
}

workflow.path = file

return workflow, nil
}
78 changes: 78 additions & 0 deletions v2/pkg/workflows/generator.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
package workflows

import (
"errors"
"strings"

"github.com/segmentio/ksuid"
)

// generateLogicFromWorkflows generates a workflow logic using the
// yaml based workflow declaration.
//
// The implementation is very basic and contains a simple yaml->tengo
// convertor that implements basic required features.
func (w *Workflow) generateLogicFromWorkflows() error {
w.Variables = make(map[string]string)

workflowBuilder := &strings.Builder{}
for _, template := range w.Workflows {
if err := w.generateTemplateFunc(template, workflowBuilder); err != nil {
return err
}
}
w.Logic = workflowBuilder.String()
return nil
}

func (w *Workflow) generateTemplateFunc(template *WorkflowTemplate, workflowBuilder *strings.Builder) error {
builder := &strings.Builder{}
builder.WriteString("var_")
builder.WriteString(ksuid.New().String())
ID := builder.String()
w.Variables[ID] = template.Template

if len(template.Subtemplates) > 0 && len(template.Matchers) > 0 {
return errors.New("subtemplates and matchers cannot be present together")
}
workflowBuilder.WriteRune('\n')
if len(template.Matchers) > 0 {
workflowBuilder.WriteString(ID)
workflowBuilder.WriteString("()\n")

for _, matcher := range template.Matchers {
if len(matcher.Subtemplates) == 0 {
return errors.New("no subtemplates present for matcher")
}
workflowBuilder.WriteString("\nif ")
workflowBuilder.WriteString(ID)
workflowBuilder.WriteString("[\"")
workflowBuilder.WriteString(matcher.Name)
workflowBuilder.WriteString("\"] {")

for _, subtemplate := range matcher.Subtemplates {
if err := w.generateTemplateFunc(subtemplate, workflowBuilder); err != nil {
return err
}
}
workflowBuilder.WriteString("\n}")
}
}
if len(template.Subtemplates) > 0 {
workflowBuilder.WriteString("if ")
workflowBuilder.WriteString(ID)
workflowBuilder.WriteString("() {")

for _, subtemplate := range template.Subtemplates {
if err := w.generateTemplateFunc(subtemplate, workflowBuilder); err != nil {
return err
}
}
workflowBuilder.WriteString("\n}")
}
if len(template.Matchers) == 0 && len(template.Subtemplates) == 0 {
workflowBuilder.WriteString(ID)
workflowBuilder.WriteString("();")
}
return nil
}
17 changes: 16 additions & 1 deletion v2/pkg/workflows/workflows.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,22 @@ type Workflow struct {
Variables map[string]string `yaml:"variables"`
// Logic contains the workflow pseudo-code
Logic string `yaml:"logic"`
path string
// Workflows is a yaml based workflow declaration code.
Workflows []*WorkflowTemplate `yaml:"workflows"`
path string
}

// WorkflowTemplate is a template to be ran as part of a workflow
type WorkflowTemplate struct {
Template string `yaml:"template"`
Matchers []*Matcher `yaml:"matchers"`
Subtemplates []*WorkflowTemplate `yaml:"subtemplates"`
}

// Matcher performs conditional matching on the workflow template results.
type Matcher struct {
Name string `yaml:"name"`
Subtemplates []*WorkflowTemplate `yaml:"subtemplates"`
}

// GetPath of the workflow
Expand Down

0 comments on commit d4e03a0

Please sign in to comment.