Merge pull request #20 from projectdiscovery/maintenance-code-refactor

code refactor

Author: Mzack9999

.github/workflows/build.yaml

@@ -32,8 +32,8 @@ jobs:
 
       - name: Test
         run: go test .
-        working-directory: .
+        working-directory: cmd/simplehttpserver/
 
       - name: Build
         run: go build .
-        working-directory: .
+        working-directory: cmd/simplehttpserver/

.goreleaser.yml

@@ -1,6 +1,6 @@
 builds:
     - binary: simplehttpserver
-      main: simplehttpserver.go
+      main: cmd/simplehttpserver/simplehttpserver.go
       goos:
         - linux
         - windows

Dockerfile

@@ -1,6 +1,6 @@
-FROM golang:1.14-alpine AS builder
+FROM golang:1.16-alpine AS builder
 RUN apk add --no-cache git
-RUN GO111MODULE=auto go get -u -v github.com/projectdiscovery/simplehttpserver
+RUN GO111MODULE=auto go get -u -v github.com/projectdiscovery/simplehttpserver/cmd/simplehttpserver
 
 FROM alpine:latest
 COPY --from=builder /go/bin/simplehttpserver /usr/local/bin/

gen_cert.sh

@@ -0,0 +1,24 @@
+package runner
+
+import "github.com/projectdiscovery/gologger"
+
+const banner = `
+     _                 _      _     _   _                                      
+ ___(_)_ __ ___  _ __ | | ___| |__ | |_| |_ _ __  ___  ___ _ ____   _____ _ __ 
+/ __| | '_ ' _ \| '_ \| |/ _ \ '_ \| __| __| '_ \/ __|/ _ \ '__\ \ / / _ \ '__|
+\__ \ | | | | | | |_) | |  __/ | | | |_| |_| |_) \__ \  __/ |   \ V /  __/ |   
+|___/_|_| |_| |_| .__/|_|\___|_| |_|\__|\__| .__/|___/\___|_|    \_/ \___|_|   
+			   |_|                        |_|                                 
+`
+
+// Version is the current version
+const Version = `0.0.1`
+
+// showBanner is used to show the banner to the user
+func showBanner() {
+	gologger.Print().Msgf("%s\n", banner)
+	gologger.Print().Msgf("\t\tprojectdiscovery.io\n\n")
+
+	gologger.Print().Msgf("Use with caution. You are responsible for your actions\n")
+	gologger.Print().Msgf("Developers assume no liability and are not responsible for any misuse or damage.\n")
+}

internal/runner/banner.go

@@ -0,0 +1,2 @@
+// Package runner contains the internal logic
+package runner

internal/runner/doc.go

@@ -0,0 +1,93 @@
+package runner
+
+import (
+	"flag"
+	"os"
+	"strings"
+
+	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/gologger/levels"
+)
+
+type Options struct {
+	ListenAddress  string
+	Folder         string
+	BasicAuth      string
+	username       string
+	password       string
+	Realm          string
+	TLSCertificate string
+	TLSKey         string
+	TLSDomain      string
+	HTTPS          bool
+	Verbose        bool
+	EnableUpload   bool
+	EnableTCP      bool
+	RulesFile      string
+	TCPWithTLS     bool
+	Version        bool
+	Silent         bool
+}
+
+// ParseOptions parses the command line options for application
+func ParseOptions() *Options {
+	options := &Options{}
+	flag.StringVar(&options.ListenAddress, "listen", "0.0.0.0:8000", "Address:Port")
+	flag.BoolVar(&options.EnableTCP, "tcp", false, "TCP Server")
+	flag.BoolVar(&options.TCPWithTLS, "tls", false, "Enable TCP TLS")
+	flag.StringVar(&options.RulesFile, "rules", "", "Rules yaml file")
+	flag.StringVar(&options.Folder, "path", ".", "Folder")
+	flag.BoolVar(&options.EnableUpload, "upload", false, "Enable upload via PUT")
+	flag.BoolVar(&options.HTTPS, "https", false, "HTTPS")
+	flag.StringVar(&options.TLSCertificate, "cert", "", "HTTPS Certificate")
+	flag.StringVar(&options.TLSKey, "key", "", "HTTPS Certificate Key")
+	flag.StringVar(&options.TLSDomain, "domain", "local.host", "Domain")
+	flag.BoolVar(&options.Verbose, "verbose", false, "Verbose")
+	flag.StringVar(&options.BasicAuth, "basic-auth", "", "Basic auth (username:password)")
+	flag.StringVar(&options.Realm, "realm", "Please enter username and password", "Realm")
+	flag.BoolVar(&options.Version, "version", false, "Show version of the software")
+	flag.BoolVar(&options.Silent, "silent", false, "Show only results in the output")
+
+	flag.Parse()
+
+	// Read the inputs and configure the logging
+	options.configureOutput()
+
+	showBanner()
+
+	if options.Version {
+		gologger.Info().Msgf("Current Version: %s\n", Version)
+		os.Exit(0)
+	}
+
+	options.validateOptions()
+
+	return options
+}
+
+func (options *Options) validateOptions() {
+	if flag.NArg() > 0 && options.Folder == "." {
+		options.Folder = flag.Args()[0]
+	}
+
+	if options.BasicAuth != "" {
+		baTokens := strings.SplitN(options.BasicAuth, ":", 2)
+		if len(baTokens) > 0 {
+			options.username = baTokens[0]
+		}
+		if len(baTokens) > 1 {
+			options.password = baTokens[1]
+		}
+	}
+}
+
+// configureOutput configures the output on the screen
+func (options *Options) configureOutput() {
+	// If the user desires verbose output, show verbose output
+	if options.Verbose {
+		gologger.DefaultLogger.SetMaxLevel(levels.LevelVerbose)
+	}
+	if options.Silent {
+		gologger.DefaultLogger.SetMaxLevel(levels.LevelSilent)
+	}
+}

internal/runner/options.go

@@ -0,0 +1,70 @@
+package runner
+
+import (
+	"github.com/projectdiscovery/simplehttpserver/pkg/httpserver"
+	"github.com/projectdiscovery/simplehttpserver/pkg/tcpserver"
+)
+
+// Runner is a client for running the enumeration process.
+type Runner struct {
+	options    *Options
+	serverTCP  *tcpserver.TCPServer
+	httpServer *httpserver.HTTPServer
+}
+
+func New(options *Options) (*Runner, error) {
+	r := Runner{options: options}
+	if r.options.EnableTCP {
+		serverTCP, err := tcpserver.New(tcpserver.Options{
+			Listen:  r.options.ListenAddress,
+			TLS:     r.options.TCPWithTLS,
+			Domain:  "local.host",
+			Verbose: r.options.Verbose,
+		})
+		if err != nil {
+			return nil, err
+		}
+		err = serverTCP.LoadTemplate(r.options.RulesFile)
+		if err != nil {
+			return nil, err
+		}
+		r.serverTCP = serverTCP
+		return &r, nil
+	}
+
+	httpServer, err := httpserver.New(&httpserver.Options{
+		Folder:            r.options.Folder,
+		EnableUpload:      r.options.EnableUpload,
+		ListenAddress:     r.options.ListenAddress,
+		TLS:               r.options.HTTPS,
+		Certificate:       r.options.TLSCertificate,
+		CertificateKey:    r.options.TLSKey,
+		CertificateDomain: r.options.TLSDomain,
+		BasicAuthUsername: r.options.username,
+		BasicAuthPassword: r.options.password,
+		BasicAuthReal:     r.options.Realm,
+		Verbose:           r.options.Verbose,
+	})
+	if err != nil {
+		return nil, err
+	}
+	r.httpServer = httpServer
+
+	return &r, nil
+}
+
+func (r *Runner) Run() error {
+	if r.options.EnableTCP {
+		return r.serverTCP.ListenAndServe()
+	}
+
+	if r.options.HTTPS {
+		return r.httpServer.ListenAndServeTLS()
+	}
+
+	return r.httpServer.ListenAndServe()
+}
+
+func (r *Runner) Close() error {
+	return nil
+}

internal/runner/runner.go

@@ -0,0 +1,19 @@
+package httpserver
+
+import (
+	"fmt"
+	"net/http"
+)
+
+func (t *HTTPServer) basicauthlayer(handler http.Handler) http.HandlerFunc {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		user, pass, ok := r.BasicAuth()
+		if !ok || user != t.options.BasicAuthUsername || pass != t.options.BasicAuthPassword {
+			w.Header().Set("WWW-Authenticate", fmt.Sprintf("Basic realm=\"%s\"", t.options.BasicAuthReal))
+			w.WriteHeader(http.StatusUnauthorized)
+			w.Write([]byte("Unauthorized.\n")) //nolint
+			return
+		}
+		handler.ServeHTTP(w, r)
+	})
+}

pkg/httpserver/authlayer.go

@@ -0,0 +1,120 @@
+package httpserver
+
+import (
+	"errors"
+	"net"
+	"net/http"
+	"os"
+	"runtime"
+	"strconv"
+	"syscall"
+
+	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/simplehttpserver/pkg/sslcert"
+)
+
+type Options struct {
+	Folder            string
+	EnableUpload      bool
+	ListenAddress     string
+	TLS               bool
+	Certificate       string
+	CertificateKey    string
+	CertificateDomain string
+	BasicAuthUsername string
+	BasicAuthPassword string
+	BasicAuthReal     string
+	Verbose           bool
+}
+
+type HTTPServer struct {
+	options  *Options
+	layers   http.Handler
+	listener net.Listener
+}
+
+func New(options *Options) (*HTTPServer, error) {
+	var h HTTPServer
+	EnableUpload = options.EnableUpload
+	EnableVerbose = options.Verbose
+	layers := h.loglayer(http.FileServer(http.Dir(options.Folder)))
+	if options.BasicAuthUsername != "" || options.BasicAuthPassword != "" {
+		layers = h.loglayer(h.basicauthlayer(http.FileServer(http.Dir(options.Folder))))
+	}
+
+	return &HTTPServer{options: options, layers: layers}, nil
+}
+
+func (t *HTTPServer) ListenAndServe() error {
+	var err error
+retry_listen:
+	gologger.Print().Msgf("Serving %s on http://%s/...", t.options.Folder, t.options.ListenAddress)
+	err = http.ListenAndServe(t.options.ListenAddress, t.layers)
+	if err != nil {
+		if isErrorAddressAlreadyInUse(err) {
+			gologger.Print().Msgf("Can't listen on %s: %s - retrying with another port\n", t.options.ListenAddress, err)
+			newListenAddress, err := incPort(t.options.ListenAddress)
+			if err != nil {
+				return err
+			}
+			t.options.ListenAddress = newListenAddress
+			goto retry_listen
+		}
+	}
+	return nil
+}
+
+func (t *HTTPServer) ListenAndServeTLS() error {
+	gologger.Print().Msgf("Serving %s on https://%s/...", t.options.Folder, t.options.ListenAddress)
+	if t.options.Certificate == "" || t.options.CertificateKey == "" {
+		tlsOptions := sslcert.DefaultOptions
+		tlsOptions.Host = t.options.CertificateDomain
+		tlsConfig, err := sslcert.NewTLSConfig(tlsOptions)
+		if err != nil {
+			return err
+		}
+		httpServer := &http.Server{
+			Addr:      t.options.ListenAddress,
+			TLSConfig: tlsConfig,
+		}
+		httpServer.Handler = t.layers
+		return httpServer.ListenAndServeTLS("", "")
+	}
+	return http.ListenAndServeTLS(t.options.ListenAddress, t.options.Certificate, t.options.CertificateKey, t.layers)
+}
+
+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
+}