From 4fb2abb558e98ae03d11f2220d7f0d2618efb0e8 Mon Sep 17 00:00:00 2001 From: c-f <35263248+c-f@users.noreply.github.com> Date: Mon, 20 Dec 2021 16:46:39 +0100 Subject: [PATCH 1/3] remove bug of reading previous messages --- pkg/tcpserver/tcpserver.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pkg/tcpserver/tcpserver.go b/pkg/tcpserver/tcpserver.go index 876fbb4..bed8e95 100644 --- a/pkg/tcpserver/tcpserver.go +++ b/pkg/tcpserver/tcpserver.go @@ -59,14 +59,14 @@ func (t *TCPServer) handleConnection(conn net.Conn) error { if err := conn.SetReadDeadline(time.Now().Add(readTimeout * time.Second)); err != nil { gologger.Info().Msgf("%s\n", err) } - _, err := conn.Read(buf) + n, err := conn.Read(buf) if err != nil { return err } - gologger.Print().Msgf("%s\n", buf) + gologger.Print().Msgf("%s\n", buf[:n]) - resp, err := t.BuildResponse(buf) + resp, err := t.BuildResponse(buf[:n]) if err != nil { return err } From e615365be58ccaf78274d0c212a35bf62a9d382a Mon Sep 17 00:00:00 2001 From: c-f <35263248+c-f@users.noreply.github.com> Date: Mon, 20 Dec 2021 17:44:21 +0100 Subject: [PATCH 2/3] add livereloading for rule config --- internal/runner/runner.go | 2 ++ internal/runner/watchdog.go | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+) create mode 100644 internal/runner/watchdog.go diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 269dd15..cbd9227 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -42,6 +42,8 @@ func New(options *Options) (*Runner, error) { if err != nil { return nil, err } + watchFile(r.options.RulesFile, serverTCP.LoadTemplate) + r.serverTCP = serverTCP return &r, nil } diff --git a/internal/runner/watchdog.go b/internal/runner/watchdog.go new file mode 100644 index 0000000..2cdde4c --- /dev/null +++ b/internal/runner/watchdog.go @@ -0,0 +1,36 @@ +package runner + +import ( + "log" + + "github.com/fsnotify/fsnotify" +) + +type WatchEvent func(fname string) error + +func watchFile(fname string, callback WatchEvent) (watcher *fsnotify.Watcher, err error) { + watcher, err = fsnotify.NewWatcher() + if err != nil { + return + } + go func() { + for { + select { + case event, ok := <-watcher.Events: + if !ok { + continue + } + if event.Op&fsnotify.Write == fsnotify.Write { + if err := callback(fname); err != nil { + log.Println("err", err) + } + } + case <-watcher.Errors: + // ignore errors for now + } + } + }() + + err = watcher.Add(fname) + return +} From 420a99f54aea3a31ce03055cdc0c7792440b9481 Mon Sep 17 00:00:00 2001 From: c-f <35263248+c-f@users.noreply.github.com> Date: Mon, 20 Dec 2021 18:11:53 +0100 Subject: [PATCH 3/3] make linter happy --- README.md | 15 ++++++- go.mod | 1 + go.sum | 4 ++ internal/runner/runner.go | 6 ++- pkg/tcpserver/addr.go | 9 ++++ pkg/tcpserver/responseengine.go | 7 ++- pkg/tcpserver/rule.go | 50 ++++++++++++++++++--- pkg/tcpserver/tcpserver.go | 77 ++++++++++++++++++++++++++++++--- 8 files changed, 153 insertions(+), 16 deletions(-) create mode 100644 pkg/tcpserver/addr.go diff --git a/README.md b/README.md index b533b2a..e08d29a 100644 --- a/README.md +++ b/README.md @@ -128,7 +128,9 @@ simplehttpserver -rule rules.yaml -tcp -tls -domain localhost The rules are written as follows: ```yaml rules: - - match: regex + - match: regex-match + match-contains: literal-match + name: rule-name response: response data ``` @@ -137,6 +139,7 @@ For example to handle two different paths simulating an HTTP server or SMTP comm rules: # HTTP Requests - match: GET /path1 + name: redirect response: | HTTP/1.0 200 OK Server: httpd/2.0 @@ -149,6 +152,7 @@ rules: - match: GET /path2 + name: "404" response: | HTTP/1.0 404 OK Server: httpd/2.0 @@ -156,6 +160,7 @@ rules: Not found # SMTP Commands - match: "EHLO example.com" + name: smtp response: | 250-localhost Nice to meet you, [127.0.0.1] 250-PIPELINING @@ -167,6 +172,14 @@ rules: response: 250 Accepted - match: "RCPT TO: " response: 250 Accepted + + - match-contains: !!binary | + MAwCAQFgBwIBAwQAgAA= + name: "ldap" + # Request: 300c 0201 0160 0702 0103 0400 8000 0....`........ + # Response: 300c 0201 0161 070a 0100 0400 0400 0....a........ + response: !!binary | + MAwCAQFhBwoBAAQABAA= ``` ## Note diff --git a/go.mod b/go.mod index f4c4a9d..7b6de7d 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/projectdiscovery/simplehttpserver go 1.14 require ( + github.com/fsnotify/fsnotify v1.5.1 // indirect github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2 github.com/projectdiscovery/gologger v1.1.4 github.com/projectdiscovery/sslcert v0.0.0-20210416140253-8f56bec1bb5e diff --git a/go.sum b/go.sum index aafe4a2..e239b8e 100644 --- a/go.sum +++ b/go.sum @@ -2,6 +2,8 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/fsnotify/fsnotify v1.5.1 h1:mZcQUHVQUQWoPXXtuf9yuEXKudkV2sx1E06UadKWpgI= +github.com/fsnotify/fsnotify v1.5.1/go.mod h1:T3375wBYaZdLLcVNkcVbzGHY7f1l/uK5T5Ai1i3InKU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/json-iterator/go v1.1.10 h1:Kz6Cvnvv2wGdaG/V8yMvfkmNiXq9Ya2KUv4rouJJr68= github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= @@ -31,6 +33,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c h1:F1jZWGFhYfh0Ci55sIpILtKKK8p3i2/krTr0H1rg74I= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= diff --git a/internal/runner/runner.go b/internal/runner/runner.go index cbd9227..49d5cd3 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -42,7 +42,11 @@ func New(options *Options) (*Runner, error) { if err != nil { return nil, err } - watchFile(r.options.RulesFile, serverTCP.LoadTemplate) + watcher, err := watchFile(r.options.RulesFile, serverTCP.LoadTemplate) + if err != nil { + return nil, err + } + defer watcher.Close() r.serverTCP = serverTCP return &r, nil diff --git a/pkg/tcpserver/addr.go b/pkg/tcpserver/addr.go new file mode 100644 index 0000000..b678b30 --- /dev/null +++ b/pkg/tcpserver/addr.go @@ -0,0 +1,9 @@ +package tcpserver + +// ContextType is the key type stored in ctx +type ContextType string + +var ( + // Addr is the contextKey where the net.Addr is stored + Addr ContextType = "addr" +) diff --git a/pkg/tcpserver/responseengine.go b/pkg/tcpserver/responseengine.go index ec15da0..80fb795 100644 --- a/pkg/tcpserver/responseengine.go +++ b/pkg/tcpserver/responseengine.go @@ -6,9 +6,12 @@ import ( // BuildResponse according to rules func (t *TCPServer) BuildResponse(data []byte) ([]byte, error) { + t.mux.RLock() + defer t.mux.RUnlock() + // Process all the rules - for _, rule := range t.options.rules { - if rule.matchRegex.Match(data) { + for _, rule := range t.rules { + if rule.MatchInput(data) { return []byte(rule.Response), nil } } diff --git a/pkg/tcpserver/rule.go b/pkg/tcpserver/rule.go index 903331b..aa9e6e8 100644 --- a/pkg/tcpserver/rule.go +++ b/pkg/tcpserver/rule.go @@ -1,6 +1,9 @@ package tcpserver -import "regexp" +import ( + "regexp" + "strings" +) // RulesConfiguration from yaml type RulesConfiguration struct { @@ -9,13 +12,20 @@ type RulesConfiguration struct { // Rule to apply to various requests type Rule struct { - Match string `yaml:"match,omitempty"` - matchRegex *regexp.Regexp - Response string `yaml:"response,omitempty"` + Name string `yaml:"name,omitempty"` + Match string `yaml:"match,omitempty"` + MatchContains string `yaml:"match-contains,omitempty"` + matchRegex *regexp.Regexp + Response string `yaml:"response,omitempty"` } -// NewRule from model +// NewRule creates a new Rule - default is regex func NewRule(match, response string) (*Rule, error) { + return NewRegexRule(match, response) +} + +// NewRegexRule returns a new regex-match Rule +func NewRegexRule(match, response string) (*Rule, error) { regxp, err := regexp.Compile(match) if err != nil { return nil, err @@ -23,3 +33,33 @@ func NewRule(match, response string) (*Rule, error) { return &Rule{Match: match, matchRegex: regxp, Response: response}, nil } + +// NewLiteralRule returns a new literal-match Rule +func NewLiteralRule(match, response string) (*Rule, error) { + return &Rule{MatchContains: match, Response: response}, nil +} + +// NewRuleFromTemplate "copies" a new Rule +func NewRuleFromTemplate(r Rule) (newRule *Rule, err error) { + newRule = &Rule{ + Name: r.Name, + Response: r.Response, + MatchContains: r.MatchContains, + Match: r.Match, + } + if newRule.Match != "" { + newRule.matchRegex, err = regexp.Compile(newRule.Match) + } + + return +} + +// MatchInput returns if the input was matches with one of the matchers +func (r *Rule) MatchInput(input []byte) bool { + if r.matchRegex != nil && r.matchRegex.Match(input) { + return true + } else if r.MatchContains != "" && strings.Contains(string(input), r.MatchContains) { + return true + } + return false +} diff --git a/pkg/tcpserver/tcpserver.go b/pkg/tcpserver/tcpserver.go index bed8e95..cbdd407 100644 --- a/pkg/tcpserver/tcpserver.go +++ b/pkg/tcpserver/tcpserver.go @@ -1,9 +1,12 @@ package tcpserver import ( + "context" "crypto/tls" + "errors" "io/ioutil" "net" + "sync" "time" "github.com/projectdiscovery/gologger" @@ -24,20 +27,35 @@ type Options struct { Verbose bool } +// CallBackFunc handles what is send back to the client, based on the incomming question +type CallBackFunc func(ctx context.Context, question []byte) (answer []byte, err error) + // TCPServer instance type TCPServer struct { options *Options listener net.Listener + + // Callbacks to retrieve information about the system + HandleMessageFnc CallBackFunc + + mux sync.RWMutex + rules []Rule } // New tcp server instance with specified options func New(options *Options) (*TCPServer, error) { - return &TCPServer{options: options}, nil + srv := &TCPServer{options: options} + srv.HandleMessageFnc = srv.BuildResponseWithContext + srv.rules = options.rules + return srv, nil } // AddRule to the server func (t *TCPServer) AddRule(rule Rule) error { - t.options.rules = append(t.options.rules, rule) + t.mux.Lock() + defer t.mux.Unlock() + + t.rules = append(t.rules, rule) return nil } @@ -51,9 +69,12 @@ func (t *TCPServer) ListenAndServe() error { return t.run() } -func (t *TCPServer) handleConnection(conn net.Conn) error { +func (t *TCPServer) handleConnection(conn net.Conn, callback CallBackFunc) error { defer conn.Close() //nolint + // Create Context + ctx := context.WithValue(context.Background(), Addr, conn.RemoteAddr()) + buf := make([]byte, 4096) for { if err := conn.SetReadDeadline(time.Now().Add(readTimeout * time.Second)); err != nil { @@ -66,8 +87,9 @@ func (t *TCPServer) handleConnection(conn net.Conn) error { gologger.Print().Msgf("%s\n", buf[:n]) - resp, err := t.BuildResponse(buf[:n]) + resp, err := callback(ctx, buf[:n]) if err != nil { + gologger.Info().Msgf("Closing connection: %s\n", err) return err } @@ -112,7 +134,7 @@ func (t *TCPServer) run() error { if err != nil { return err } - go t.handleConnection(c) //nolint + go t.handleConnection(c, t.HandleMessageFnc) //nolint } } @@ -133,13 +155,54 @@ func (t *TCPServer) LoadTemplate(templatePath string) error { return err } + t.mux.Lock() + defer t.mux.Unlock() + + t.rules = make([]Rule, 0) for _, ruleTemplate := range config.Rules { - rule, err := NewRule(ruleTemplate.Match, ruleTemplate.Response) + rule, err := NewRuleFromTemplate(ruleTemplate) if err != nil { return err } - t.options.rules = append(t.options.rules, *rule) + t.rules = append(t.rules, *rule) } + gologger.Info().Msgf("TCP configuration loaded. Rules: %d\n", len(t.rules)) + return nil } + +// MatchRule returns the rule, which was matched first +func (t *TCPServer) MatchRule(data []byte) (rule Rule, err error) { + t.mux.RLock() + defer t.mux.RUnlock() + + // Process all the rules + for _, rule := range t.rules { + if rule.MatchInput(data) { + return rule, nil + } + } + return Rule{}, errors.New("no matched rule") +} + +// BuildResponseWithContext is a wrapper with context +func (t *TCPServer) BuildResponseWithContext(ctx context.Context, data []byte) ([]byte, error) { + return t.BuildResponse(data) +} + +// BuildResponseWithContext is a wrapper with context +func (t *TCPServer) BuildRuleResponse(ctx context.Context, data []byte) ([]byte, error) { + addr := "unknown" + if netAddr, ok := ctx.Value(Addr).(net.Addr); ok { + addr = netAddr.String() + } + rule, err := t.MatchRule(data) + if err != nil { + return []byte(":) "), err + } + + gologger.Info().Msgf("Incoming TCP request(%s) from: %s\n", rule.Name, addr) + + return []byte(rule.Response), nil +}