Merge branch 'dev' into feature-generating-self-signed-certificates

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,125 @@
+package tcpserver
+
+import (
+	"crypto/tls"
+	"io/ioutil"
+	"net"
+	"time"
+
+	"github.com/projectdiscovery/simplehttpserver/pkg/sslcert"
+	"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 = t.options.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

@@ -2,45 +2,58 @@ package main
 
 import (
 	"bytes"
+	"errors"
 	"flag"
 	"fmt"
 	"io/ioutil"
-	"log"
+	"net"
 	"net/http"
 	"net/http/httputil"
+	"os"
 	"path"
+	"runtime"
+	"strconv"
+	"strings"
+	"syscall"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/simplehttpserver/pkg/sslcert"
+	"github.com/projectdiscovery/simplehttpserver/pkg/tcpserver"
 )
 
 type options struct {
 	ListenAddress string
 	Folder        string
-	Username      string
-	Password      string
+	BasicAuth     string
+	username      string
+	password      string
 	Realm         string
 	Certificate   string
 	Key           string
 	Domain        string
 	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")
 	flag.StringVar(&opts.Certificate, "cert", "", "Certificate")
 	flag.StringVar(&opts.Key, "key", "", "Key")
 	flag.StringVar(&opts.Domain, "domain", "local.host", "Domain")
 	flag.BoolVar(&opts.Verbose, "v", false, "Verbose")
-	flag.StringVar(&opts.Username, "username", "", "Basic auth username")
-	flag.StringVar(&opts.Password, "password", "", "Basic auth password")
+	flag.StringVar(&opts.BasicAuth, "basic-auth", "", "Basic auth (username:password)")
 	flag.StringVar(&opts.Realm, "realm", "Please enter username and password", "Realm")
 
 	flag.Parse()
@@ -49,15 +62,37 @@ 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.Username != "" || opts.Password != "" {
+	if opts.BasicAuth != "" {
+		baTokens := strings.SplitN(opts.BasicAuth, ":", 2)
+		if len(baTokens) > 0 {
+			opts.username = baTokens[0]
+		}
+		if len(baTokens) > 1 {
+			opts.password = baTokens[1]
+		}
 		layers = loglayer(basicauthlayer(http.FileServer(http.Dir(opts.Folder))))
 	}
-
 	if opts.Upload {
-		gologger.Print().Msgf("Upload enabled")
+		gologger.Print().Msg("Starting service with Upload enabled")
 	}
+retry_listen:
+	gologger.Print().Msgf("Serving %s on http://%s/...", opts.Folder, opts.ListenAddress)
+	var err error
 	if opts.HTTPS {
 		if opts.Certificate == "" || opts.Key == "" {
 			tlsOptions := sslcert.DefaultOptions
@@ -76,7 +111,19 @@ func main() {
 			gologger.Print().Msgf("%s\n", http.ListenAndServeTLS(opts.ListenAddress, opts.Certificate, opts.Key, layers))
 		}
 	} else {
-		gologger.Print().Msgf("%s\n", http.ListenAndServe(opts.ListenAddress, layers))
+		err = http.ListenAndServe(opts.ListenAddress, layers)
+	}
+	if err != nil {
+		if isErrorAddressAlreadyInUse(err) {
+			gologger.Print().Msgf("Can't listen on %s: %s - retrying with another port\n", opts.ListenAddress, err)
+			newListenAddress, err := incPort(opts.ListenAddress)
+			if err != nil {
+				gologger.Fatal().Msgf("%s\n", err)
+			}
+			opts.ListenAddress = newListenAddress
+			goto retry_listen
+		}
+		gologger.Print().Msgf("%s\n", err)
 	}
 }
 
@@ -90,11 +137,11 @@ func loglayer(handler http.Handler) http.Handler {
 		if opts.Upload && r.Method == http.MethodPut {
 			data, err := ioutil.ReadAll(r.Body)
 			if err != nil {
-				log.Println(err)
+				gologger.Print().Msgf("%s\n", err)
 			}
 			err = handleUpload(path.Base(r.URL.Path), data)
 			if err != nil {
-				log.Println(err)
+				gologger.Print().Msgf("%s\n", err)
 			}
 		}
 
@@ -111,7 +158,7 @@ func loglayer(handler http.Handler) http.Handler {
 func basicauthlayer(handler http.Handler) http.HandlerFunc {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		user, pass, ok := r.BasicAuth()
-		if !ok || user != opts.Username || pass != opts.Password {
+		if !ok || user != opts.username || pass != opts.password {
 			w.Header().Set("WWW-Authenticate", fmt.Sprintf("Basic realm=\"%s\"", opts.Realm))
 			w.WriteHeader(http.StatusUnauthorized)
 			w.Write([]byte("Unauthorized.\n")) //nolint
@@ -148,3 +195,39 @@ func (lrw *loggingResponseWriter) WriteHeader(code int) {
 func handleUpload(file string, data []byte) error {
 	return ioutil.WriteFile(file, data, 0655)
 }
+
+func isErrorAddressAlreadyInUse(err error) bool {
+	var eOsSyscall *os.SyscallError
+	if !errors.As(err, &eOsSyscall) {
+		return false
+	}
+	var errErrno syscall.Errno // doesn't need a "*" (ptr) because it's already a ptr (uintptr)
+	if !errors.As(eOsSyscall, &errErrno) {
+		return false
+	}
+	if errErrno == syscall.EADDRINUSE {
+		return true
+	}
+	const WSAEADDRINUSE = 10048
+	if runtime.GOOS == "windows" && errErrno == WSAEADDRINUSE {
+		return true
+	}
+	return false
+}
+
+func incPort(address string) (string, error) {
+	addrOrig, portOrig, err := net.SplitHostPort(address)
+	if err != nil {
+		return address, err
+	}
+
+	// increment port
+	portNumber, err := strconv.Atoi(portOrig)
+	if err != nil {
+		return address, err
+	}
+	portNumber++
+	newPort := strconv.FormatInt(int64(portNumber), 10)
+
+	return net.JoinHostPort(addrOrig, newPort), nil
+}