fix(schema): generation of missing JSON schema definitions (#4995)

* fix(schema): generation of missing JSON schema definitions * make headers and data to accept multi-type inputs * misc update
projectdiscovery · Apr 7, 2024 · 8c27ca2 · 8c27ca2
1 parent 82e25f6
commit 8c27ca2
Show file tree

Hide file tree

Showing 18 changed files with 119 additions and 22 deletions.
diff --git a/cmd/docgen/docgen.go b/cmd/docgen/docgen.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"log"
 	"os"
+	"reflect"
 	"regexp"
 	"strings"
 
@@ -33,6 +34,12 @@ func main() {
 
 	// Generate jsonschema
 	r := &jsonschema.Reflector{}
+	r.Namer = func(r reflect.Type) string {
+		if r.Kind() == reflect.Slice {
+			return ""
+		}
+		return r.String()
+	}
 	jsonschemaData := r.Reflect(&templates.Template{})
 
 	var buf bytes.Buffer

diff --git a/pkg/fuzz/type.go b/pkg/fuzz/type.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/invopop/jsonschema"
 	mapsutil "github.com/projectdiscovery/utils/maps"
 	"gopkg.in/yaml.v2"
 )
@@ -29,6 +30,44 @@ type SliceOrMapSlice struct {
 	KV    *mapsutil.OrderedMap[string, string]
 }
 
+func (v SliceOrMapSlice) JSONSchemaExtend(schema *jsonschema.Schema) *jsonschema.Schema {
+	schema = &jsonschema.Schema{
+		Title:       schema.Title,
+		Description: schema.Description,
+		Type:        "array",
+		Items: &jsonschema.Schema{
+			OneOf: []*jsonschema.Schema{
+				{
+					Type: "string",
+				},
+				{
+					Type: "object",
+				},
+			},
+		},
+	}
+	return schema
+}
+
+func (v SliceOrMapSlice) JSONSchema() *jsonschema.Schema {
+	gotType := &jsonschema.Schema{
+		Title:       "Payloads of Fuzz Rule",
+		Description: "Payloads to perform fuzzing substitutions with.",
+		Type:        "array",
+		Items: &jsonschema.Schema{
+			OneOf: []*jsonschema.Schema{
+				{
+					Type: "string",
+				},
+				{
+					Type: "object",
+				},
+			},
+		},
+	}
+	return gotType
+}
+
 // UnmarshalJSON implements json.Unmarshaler interface.
 func (v *SliceOrMapSlice) UnmarshalJSON(data []byte) error {
 	// try to unmashal as a string and fallback to map

diff --git a/pkg/model/types/userAgent/user_agent.go b/pkg/model/types/userAgent/user_agent.go
@@ -61,7 +61,7 @@ type UserAgentHolder struct {
 	Value UserAgent `mapping:"true"`
 }
 
-func (userAgentHolder UserAgentHolder) JSONSchemaType() *jsonschema.Schema {
+func (userAgentHolder UserAgentHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "userAgent for the headless",

diff --git a/pkg/operators/extractors/extractor_types.go b/pkg/operators/extractors/extractor_types.go
@@ -72,7 +72,7 @@ type ExtractorTypeHolder struct {
 	ExtractorType ExtractorType `mapping:"true"`
 }
 
-func (holder ExtractorTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder ExtractorTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of the extractor",

diff --git a/pkg/operators/matchers/matchers_types.go b/pkg/operators/matchers/matchers_types.go
@@ -82,7 +82,7 @@ func (t MatcherTypeHolder) String() string {
 	return t.MatcherType.String()
 }
 
-func (holder MatcherTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder MatcherTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of the matcher",

diff --git a/pkg/protocols/code/code.go b/pkg/protocols/code/code.go
@@ -48,25 +48,25 @@ var (
 type Request struct {
 	// Operators for the current request go here.
 	operators.Operators `yaml:",inline,omitempty"`
-	CompiledOperators   *operators.Operators `yaml:"-"`
+	CompiledOperators   *operators.Operators `yaml:"-" json:"-"`
 
 	// ID is the optional id of the request
 	ID string `yaml:"id,omitempty" json:"id,omitempty" jsonschema:"title=id of the request,description=ID is the optional ID of the Request"`
 	// description: |
 	//   Engine type
-	Engine []string `yaml:"engine,omitempty" jsonschema:"title=engine,description=Engine"`
+	Engine []string `yaml:"engine,omitempty" json:"engine,omitempty" jsonschema:"title=engine,description=Engine"`
 	// description: |
 	//   PreCondition is a condition which is evaluated before sending the request.
 	PreCondition string `yaml:"pre-condition,omitempty" json:"pre-condition,omitempty" jsonschema:"title=pre-condition for the request,description=PreCondition is a condition which is evaluated before sending the request"`
 	// description: |
 	//   Engine Arguments
-	Args []string `yaml:"args,omitempty" jsonschema:"title=args,description=Args"`
+	Args []string `yaml:"args,omitempty" json:"args,omitempty" jsonschema:"title=args,description=Args"`
 	// description: |
 	//   Pattern preferred for file name
-	Pattern string `yaml:"pattern,omitempty" jsonschema:"title=pattern,description=Pattern"`
+	Pattern string `yaml:"pattern,omitempty" json:"pattern,omitempty" jsonschema:"title=pattern,description=Pattern"`
 	// description: |
 	//   Source File/Snippet
-	Source string `yaml:"source,omitempty" jsonschema:"title=source file/snippet,description=Source snippet"`
+	Source string `yaml:"source,omitempty" json:"source,omitempty" jsonschema:"title=source file/snippet,description=Source snippet"`
 
 	options              *protocols.ExecutorOptions `yaml:"-" json:"-"`
 	preConditionCompiled *goja.Program              `yaml:"-" json:"-"`

diff --git a/pkg/protocols/common/generators/attack_types.go b/pkg/protocols/common/generators/attack_types.go
@@ -61,7 +61,7 @@ type AttackTypeHolder struct {
 	Value AttackType `mapping:"true"`
 }
 
-func (holder AttackTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder AttackTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of the attack",

diff --git a/pkg/protocols/common/variables/variables.go b/pkg/protocols/common/variables/variables.go
@@ -21,7 +21,7 @@ type Variable struct {
 	utils.InsertionOrderedStringMap `yaml:"-" json:"-"`
 }
 
-func (variables Variable) JSONSchemaType() *jsonschema.Schema {
+func (variables Variable) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:                 "object",
 		Title:                "variables for the request",

diff --git a/pkg/protocols/dns/dns.go b/pkg/protocols/dns/dns.go
@@ -60,7 +60,7 @@ type Request struct {
 	// examples:
 	//   - name: Use a retry of 100 to 150 generally
 	//     value: 100
-	TraceMaxRecursion int `yaml:"trace-max-recursion,omitempty"  jsonschema:"title=trace-max-recursion level for dns request,description=TraceMaxRecursion is the number of max recursion allowed for trace operations"`
+	TraceMaxRecursion int `yaml:"trace-max-recursion,omitempty" json:"trace-max-recursion,omitempty"  jsonschema:"title=trace-max-recursion level for dns request,description=TraceMaxRecursion is the number of max recursion allowed for trace operations"`
 
 	// description: |
 	//   Attack is the type of payload combinations to perform.
@@ -83,7 +83,7 @@ type Request struct {
 	Threads   int `yaml:"threads,omitempty" json:"threads,omitempty" jsonschema:"title=threads for sending requests,description=Threads specifies number of threads to use sending requests. This enables Connection Pooling"`
 	generator *generators.PayloadGenerator
 
-	CompiledOperators *operators.Operators `yaml:"-"`
+	CompiledOperators *operators.Operators `yaml:"-" json:"-"`
 	dnsClient         *retryabledns.Client
 	options           *protocols.ExecutorOptions
 

diff --git a/pkg/protocols/dns/dns_types.go b/pkg/protocols/dns/dns_types.go
@@ -92,7 +92,7 @@ func (holder DNSRequestTypeHolder) String() string {
 	return holder.DNSRequestType.String()
 }
 
-func (holder DNSRequestTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder DNSRequestTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of DNS request to make",

diff --git a/pkg/protocols/headless/engine/action.go b/pkg/protocols/headless/engine/action.go
@@ -1,6 +1,10 @@
 package engine
 
-import "strings"
+import (
+	"strings"
+
+	"github.com/invopop/jsonschema"
+)
 
 // Action is an action taken by the browser to reach a navigation
 //
@@ -29,6 +33,29 @@ type Action struct {
 	ActionType ActionTypeHolder `yaml:"action" json:"action" jsonschema:"title=action to perform,description=Type of actions to perform,enum=navigate,enum=script,enum=click,enum=rightclick,enum=text,enum=screenshot,enum=time,enum=select,enum=files,enum=waitload,enum=getresource,enum=extract,enum=setmethod,enum=addheader,enum=setheader,enum=deleteheader,enum=setbody,enum=waitevent,enum=keyboard,enum=debug,enum=sleep"`
 }
 
+func (a Action) JSONSchemaExtend(schema *jsonschema.Schema) {
+	argsSchema, ok := schema.Properties.Get("args")
+	if !ok {
+		return
+	}
+	argsSchema.PatternProperties = map[string]*jsonschema.Schema{
+		".*": {
+			OneOf: []*jsonschema.Schema{
+				{
+					Type: "string",
+				},
+				{
+					Type: "integer",
+				},
+				{
+					Type: "boolean",
+				},
+			},
+		},
+	}
+	argsSchema.Ref = ""
+}
+
 // String returns the string representation of an action
 func (a *Action) String() string {
 	builder := &strings.Builder{}

diff --git a/pkg/protocols/headless/engine/action_types.go b/pkg/protocols/headless/engine/action_types.go
@@ -171,7 +171,7 @@ type ActionTypeHolder struct {
 func (holder ActionTypeHolder) String() string {
 	return holder.ActionType.String()
 }
-func (holder ActionTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder ActionTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "action to perform",

diff --git a/pkg/protocols/http/http.go b/pkg/protocols/http/http.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/invopop/jsonschema"
 	json "github.com/json-iterator/go"
 	"github.com/pkg/errors"
 
@@ -219,6 +220,29 @@ type Request struct {
 	fuzzPreConditionOperator matchers.ConditionType `yaml:"-" json:"-"`
 }
 
+func (e Request) JSONSchemaExtend(schema *jsonschema.Schema) {
+	headersSchema, ok := schema.Properties.Get("headers")
+	if !ok {
+		return
+	}
+	headersSchema.PatternProperties = map[string]*jsonschema.Schema{
+		".*": {
+			OneOf: []*jsonschema.Schema{
+				{
+					Type: "string",
+				},
+				{
+					Type: "integer",
+				},
+				{
+					Type: "boolean",
+				},
+			},
+		},
+	}
+	headersSchema.Ref = ""
+}
+
 // Options returns executer options for http request
 func (r *Request) Options() *protocols.ExecutorOptions {
 	return r.options

diff --git a/pkg/protocols/http/http_method_types.go b/pkg/protocols/http/http_method_types.go
@@ -89,7 +89,7 @@ func (holder HTTPMethodTypeHolder) String() string {
 	return holder.MethodType.String()
 }
 
-func (holder HTTPMethodTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder HTTPMethodTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "method is the HTTP request method",

diff --git a/pkg/protocols/http/signature.go b/pkg/protocols/http/signature.go
@@ -51,7 +51,7 @@ type SignatureTypeHolder struct {
 	Value SignatureType
 }
 
-func (holder SignatureTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder SignatureTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of the signature",

diff --git a/pkg/protocols/network/network.go b/pkg/protocols/network/network.go
@@ -61,7 +61,7 @@ type Request struct {
 	// description: |
 	//   Port is the port to send network requests to. this acts as default port but is overriden if target/input contains
 	// non-http(s) ports like 80,8080,8081 etc
-	Port string `yaml:"port,omitempty" json:"port,omitempty" jsonschema:"title=port to send requests to,description=Port to send network requests to"`
+	Port string `yaml:"port,omitempty" json:"port,omitempty" jsonschema:"title=port to send requests to,description=Port to send network requests to,oneof_type=string;integer"`
 
 	// description:	|
 	//	ExcludePorts is the list of ports to exclude from being scanned . It is intended to be used with `Port` field and contains a list of ports which are ignored/skipped
@@ -91,7 +91,7 @@ type Request struct {
 
 	// Operators for the current request go here.
 	operators.Operators `yaml:",inline,omitempty"`
-	CompiledOperators   *operators.Operators `yaml:"-"`
+	CompiledOperators   *operators.Operators `yaml:"-" json:"-"`
 
 	generator *generators.PayloadGenerator
 	// cache any variables that may be needed for operation.
@@ -128,7 +128,7 @@ type Input struct {
 	// examples:
 	//   - value: "\"TEST\""
 	//   - value: "\"hex_decode('50494e47')\""
-	Data string `yaml:"data,omitempty" json:"data,omitempty" jsonschema:"title=data to send as input,description=Data is the data to send as the input"`
+	Data string `yaml:"data,omitempty" json:"data,omitempty" jsonschema:"title=data to send as input,description=Data is the data to send as the input,oneof_type=string;integer"`
 	// description: |
 	//   Type is the type of input specified in `data` field.
 	//

diff --git a/pkg/protocols/network/network_input_types.go b/pkg/protocols/network/network_input_types.go
@@ -66,7 +66,7 @@ func (holder NetworkInputTypeHolder) String() string {
 	return holder.NetworkInputType.String()
 }
 
-func (holder NetworkInputTypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder NetworkInputTypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type is the type of input data",

diff --git a/pkg/templates/types/types.go b/pkg/templates/types/types.go
@@ -92,7 +92,7 @@ type TypeHolder struct {
 	ProtocolType ProtocolType `mapping:"true"`
 }
 
-func (holder TypeHolder) JSONSchemaType() *jsonschema.Schema {
+func (holder TypeHolder) JSONSchema() *jsonschema.Schema {
 	gotType := &jsonschema.Schema{
 		Type:        "string",
 		Title:       "type of the protocol",