adding initial implementation of template based tcp server

Author: Mzack9999

go.mod

@@ -2,4 +2,7 @@ module github.com/projectdiscovery/simplehttpserver
 
 go 1.15
 
-require github.com/projectdiscovery/gologger v1.1.4
+require (
+	github.com/projectdiscovery/gologger v1.1.4
+	gopkg.in/yaml.v2 v2.4.0
+)

go.sum

@@ -30,6 +30,8 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 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=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/tcpserver/responseengine.go

@@ -0,0 +1,15 @@
+package tcpserver
+
+import (
+	"errors"
+)
+
+func (t *TCPServer) BuildResponse(data []byte) ([]byte, error) {
+	// Process all the rules
+	for _, rule := range t.options.rules {
+		if rule.matchRegex.Match(data) {
+			return []byte(rule.Response), nil
+		}
+	}
+	return nil, errors.New("No matched rule")
+}

pkg/tcpserver/rule.go

@@ -0,0 +1,22 @@
+package tcpserver
+
+import "regexp"
+
+type RulesConfiguration struct {
+	Rules []Rule `yaml:"rules"`
+}
+
+type Rule struct {
+	Match      string `yaml:"match,omitempty"`
+	matchRegex *regexp.Regexp
+	Response   string `yaml:"response,omitempty"`
+}
+
+func NewRule(match string, response string) (*Rule, error) {
+	regxp, err := regexp.Compile(match)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Rule{Match: match, matchRegex: regxp, Response: response}, nil
+}

pkg/tcpserver/tcpserver.go

@@ -0,0 +1,124 @@
+package tcpserver
+
+import (
+	"crypto/tls"
+	"io/ioutil"
+	"net"
+	"time"
+
+	"gopkg.in/yaml.v2"
+)
+
+type Options struct {
+	Listen      string
+	TLS         bool
+	Certificate string
+	Key         string
+	Domain      string
+	rules       []Rule
+}
+
+type TCPServer struct {
+	options  Options
+	listener net.Listener
+}
+
+func New(options Options) (*TCPServer, error) {
+	return &TCPServer{options: options}, nil
+}
+
+func (t *TCPServer) AddRule(rule Rule) error {
+	t.options.rules = append(t.options.rules, rule)
+	return nil
+}
+
+func (t *TCPServer) ListenAndServe() error {
+	listener, err := net.Listen("tcp4", t.options.Listen)
+	if err != nil {
+		return err
+	}
+	t.listener = listener
+	return t.run()
+}
+
+func (t *TCPServer) handleConnection(conn net.Conn) error {
+	defer conn.Close()
+
+	buf := make([]byte, 4096)
+	for {
+		conn.SetReadDeadline(time.Now().Add(time.Duration(5 * time.Second)))
+		_, err := conn.Read(buf)
+		if err != nil {
+			return err
+		}
+
+		resp, err := t.BuildResponse(buf)
+		if err != nil {
+			return err
+		}
+
+		conn.Write(resp)
+	}
+}
+
+func (t *TCPServer) ListenAndServeTLS() error {
+	var tlsConfig *tls.Config
+	if t.options.Certificate != "" && t.options.Key != "" {
+		cert, err := tls.LoadX509KeyPair(t.options.Certificate, t.options.Key)
+		if err != nil {
+			return err
+		}
+		tlsConfig = &tls.Config{Certificates: []tls.Certificate{cert}}
+	} else {
+		tlsOptions := sslcert.DefaultOptions
+		tlsOptions.Host = opts.Domain
+		cfg, err := sslcert.NewTLSConfig(tlsOptions)
+		if err != nil {
+			return err
+		}
+		tlsConfig = cfg
+	}
+
+	listener, err := tls.Listen("tcp", t.options.Listen, tlsConfig)
+	if err != nil {
+		return err
+	}
+	t.listener = listener
+	return t.run()
+}
+
+func (t *TCPServer) run() error {
+	for {
+		c, err := t.listener.Accept()
+		if err != nil {
+			return err
+		}
+		go t.handleConnection(c)
+	}
+}
+
+func (t *TCPServer) Close() error {
+	return t.listener.Close()
+}
+
+func (t *TCPServer) LoadTemplate(templatePath string) error {
+	var config RulesConfiguration
+	yamlFile, err := ioutil.ReadFile(templatePath)
+	if err != nil {
+		return err
+	}
+	err = yaml.Unmarshal(yamlFile, &config)
+	if err != nil {
+		return err
+	}
+
+	for _, ruleTemplate := range config.Rules {
+		rule, err := NewRule(ruleTemplate.Match, ruleTemplate.Response)
+		if err != nil {
+			return err
+		}
+		t.options.rules = append(t.options.rules, *rule)
+	}
+
+	return nil
+}

simplehttpserver.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/simplehttpserver/pkg/tcpserver"
 )
 
 type options struct {
@@ -26,12 +27,18 @@ type options struct {
 	HTTPS         bool
 	Verbose       bool
 	Upload        bool
+	TCP           bool
+	RulesFile     string
+	TLS           bool
 }
 
 var opts options
 
 func main() {
 	flag.StringVar(&opts.ListenAddress, "listen", "0.0.0.0:8000", "Address:Port")
+	flag.BoolVar(&opts.TCP, "tcp", false, "TCP Server")
+	flag.BoolVar(&opts.TLS, "tls", false, "Enable TCP TLS")
+	flag.StringVar(&opts.RulesFile, "rules", "", "Rules yaml file")
 	flag.StringVar(&opts.Folder, "path", ".", "Folder")
 	flag.BoolVar(&opts.Upload, "upload", false, "Enable upload via PUT")
 	flag.BoolVar(&opts.HTTPS, "https", false, "HTTPS")
@@ -47,6 +54,19 @@ func main() {
 		opts.Folder = flag.Args()[0]
 	}
 
+	if opts.TCP {
+		serverTCP, err := tcpserver.New(tcpserver.Options{Listen: opts.ListenAddress, TLS: opts.TLS, Domain: "local.host"})
+		if err != nil {
+			gologger.Fatal().Msgf("%s\n", err)
+		}
+		err = serverTCP.LoadTemplate(opts.RulesFile)
+		if err != nil {
+			gologger.Fatal().Msgf("%s\n", err)
+		}
+
+		gologger.Print().Msgf("%s\n", serverTCP.ListenAndServe())
+	}
+
 	gologger.Print().Msgf("Serving %s on http://%s/...", opts.Folder, opts.ListenAddress)
 	layers := loglayer(http.FileServer(http.Dir(opts.Folder)))
 	if opts.BasicAuth != "" {