mirror of https://github.com/daffainfo/nuclei.git
108 lines
2.8 KiB
Go
108 lines
2.8 KiB
Go
package main
|
|
|
|
import (
|
|
"net"
|
|
"strings"
|
|
|
|
"github.com/gobwas/ws/wsutil"
|
|
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/testutils"
|
|
)
|
|
|
|
var websocketTestCases = []TestCaseInfo{
|
|
{Path: "protocols/websocket/basic.yaml", TestCase: &websocketBasic{}},
|
|
{Path: "protocols/websocket/cswsh.yaml", TestCase: &websocketCswsh{}},
|
|
{Path: "protocols/websocket/no-cswsh.yaml", TestCase: &websocketNoCswsh{}},
|
|
{Path: "protocols/websocket/path.yaml", TestCase: &websocketWithPath{}},
|
|
}
|
|
|
|
type websocketBasic struct{}
|
|
|
|
// Execute executes a test case and returns an error if occurred
|
|
func (h *websocketBasic) Execute(filePath string) error {
|
|
connHandler := func(conn net.Conn) {
|
|
for {
|
|
msg, op, _ := wsutil.ReadClientData(conn)
|
|
if string(msg) != "hello" {
|
|
return
|
|
}
|
|
_ = wsutil.WriteServerMessage(conn, op, []byte("world"))
|
|
}
|
|
}
|
|
originValidate := func(origin string) bool {
|
|
return true
|
|
}
|
|
ts := testutils.NewWebsocketServer("", connHandler, originValidate)
|
|
defer ts.Close()
|
|
|
|
results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return expectResultsCount(results, 1)
|
|
}
|
|
|
|
type websocketCswsh struct{}
|
|
|
|
// Execute executes a test case and returns an error if occurred
|
|
func (h *websocketCswsh) Execute(filePath string) error {
|
|
connHandler := func(conn net.Conn) {
|
|
|
|
}
|
|
originValidate := func(origin string) bool {
|
|
return true
|
|
}
|
|
ts := testutils.NewWebsocketServer("", connHandler, originValidate)
|
|
defer ts.Close()
|
|
|
|
results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return expectResultsCount(results, 1)
|
|
}
|
|
|
|
type websocketNoCswsh struct{}
|
|
|
|
// Execute executes a test case and returns an error if occurred
|
|
func (h *websocketNoCswsh) Execute(filePath string) error {
|
|
connHandler := func(conn net.Conn) {
|
|
|
|
}
|
|
originValidate := func(origin string) bool {
|
|
return origin == "https://google.com"
|
|
}
|
|
ts := testutils.NewWebsocketServer("", connHandler, originValidate)
|
|
defer ts.Close()
|
|
|
|
results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return expectResultsCount(results, 0)
|
|
}
|
|
|
|
type websocketWithPath struct{}
|
|
|
|
// Execute executes a test case and returns an error if occurred
|
|
func (h *websocketWithPath) Execute(filePath string) error {
|
|
connHandler := func(conn net.Conn) {
|
|
|
|
}
|
|
originValidate := func(origin string) bool {
|
|
return origin == "https://google.com"
|
|
}
|
|
ts := testutils.NewWebsocketServer("/test", connHandler, originValidate)
|
|
defer ts.Close()
|
|
|
|
results, err := testutils.RunNucleiTemplateAndGetResults(filePath, strings.ReplaceAll(ts.URL, "http", "ws"), debug)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return expectResultsCount(results, 0)
|
|
}
|