Merge pull request #61 from c-f/livereload

ehsandeep · web-flow · commit bbfd855cbb74 · 2021-12-23T09:25:52.000+05:30
Extension of TCP Server (livereload, binary matching, custom logic)
diff --git 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,13 +152,15 @@ rules:
               <HTML><HEAD><script>top.location.href='/Main_Login.asp';</script>
               </HEAD></HTML>
   - match: GET /path2
+    name: "404"
     response: |
               HTTP/1.0 404 OK
               Server: httpd/2.0
             
               <HTML><HEAD></HEAD><BODY>Not found</BODY></HTML>
   # 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: <test@example.com>"
     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
@@ -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
@@ -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
@@ -42,6 +42,12 @@ func New(options *Options) (*Runner, error) {
 		if err != nil {
 			return nil, err
 		}
+		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/internal/runner/watchdog.go 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
+}
diff --git a/pkg/tcpserver/addr.go 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
@@ -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
@@ -1,6 +1,9 @@
 package tcpserver
 
-import "regexp"
+import (
+	"regexp"
+	"strings"
+)
 
 // RulesConfiguration from yaml
 type RulesConfiguration struct {
@@ -9,17 +12,54 @@ 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
 	}
 
 	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
@@ -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,23 +69,27 @@ 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 {
 			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 := 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
+}

Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,12 @@ func New(options Options) (Runner, error) {`
`42`	`42`	`if err != nil {`
`43`	`43`	`return nil, err`
`44`	`44`	`}`
	`45`	`+ watcher, err := watchFile(r.options.RulesFile, serverTCP.LoadTemplate)`
	`46`	`+ if err != nil {`
	`47`	`+ return nil, err`
	`48`	`+ }`
	`49`	`+ defer watcher.Close()`
	`50`	`+`
`45`	`51`	`r.serverTCP = serverTCP`
`46`	`52`	`return &r, nil`
`47`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -6,9 +6,12 @@ import (`
`6`	`6`
`7`	`7`	`// BuildResponse according to rules`
`8`	`8`	`func (t *TCPServer) BuildResponse(data []byte) ([]byte, error) {`
	`9`	`+ t.mux.RLock()`
	`10`	`+ defer t.mux.RUnlock()`
	`11`	`+`
`9`	`12`	`// Process all the rules`
`10`		`- for _, rule := range t.options.rules {`
`11`		`- if rule.matchRegex.Match(data) {`
	`13`	`+ for _, rule := range t.rules {`
	`14`	`+ if rule.MatchInput(data) {`
`12`	`15`	`return []byte(rule.Response), nil`
`13`	`16`	`}`
`14`	`17`	`}`