diff --git a/internal/runner/options.go b/internal/runner/options.go index a5869d6..95b7e03 100644 --- a/internal/runner/options.go +++ b/internal/runner/options.go @@ -12,25 +12,26 @@ import ( // Options of the tool 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 - Sandbox bool - MaxFileSize int + 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 + Sandbox bool + MaxFileSize int + MaxDumpBodySize int } // ParseOptions parses the command line options for application @@ -57,6 +58,7 @@ func ParseOptions() *Options { flag.BoolVar(&options.Silent, "silent", false, "Show only results in the output") flag.BoolVar(&options.Sandbox, "sandbox", false, "Enable sandbox mode") flag.IntVar(&options.MaxFileSize, "max-file-size", 50, "Max Upload File Size") + flag.IntVar(&options.MaxDumpBodySize, "max-dump-body-size", -1, "Max Dump Body Size") flag.Parse() diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 5806044..e9f6cbe 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -5,6 +5,7 @@ import ( "github.com/projectdiscovery/simplehttpserver/pkg/binder" "github.com/projectdiscovery/simplehttpserver/pkg/httpserver" "github.com/projectdiscovery/simplehttpserver/pkg/tcpserver" + "github.com/projectdiscovery/simplehttpserver/pkg/unit" ) // Runner is a client for running the enumeration process. @@ -59,6 +60,7 @@ func New(options *Options) (*Runner, error) { Verbose: r.options.Verbose, Sandbox: r.options.Sandbox, MaxFileSize: r.options.MaxFileSize, + MaxDumpBodySize: unit.ToMb(r.options.MaxDumpBodySize), }) if err != nil { return nil, err diff --git a/pkg/httpserver/httpserver.go b/pkg/httpserver/httpserver.go index 72da466..d1d8fef 100644 --- a/pkg/httpserver/httpserver.go +++ b/pkg/httpserver/httpserver.go @@ -24,6 +24,7 @@ type Options struct { Verbose bool Sandbox bool MaxFileSize int // 50Mb + MaxDumpBodySize int64 } // HTTPServer instance diff --git a/pkg/httpserver/loglayer.go b/pkg/httpserver/loglayer.go index 0e1a87a..04c5550 100644 --- a/pkg/httpserver/loglayer.go +++ b/pkg/httpserver/loglayer.go @@ -9,6 +9,7 @@ import ( "path/filepath" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/simplehttpserver/pkg/unit" ) // Convenience globals @@ -17,10 +18,19 @@ var ( EnableVerbose bool ) +func (t *HTTPServer) shouldDumpBody(bodysize int64) bool { + return t.options.MaxDumpBodySize > 0 && bodysize > t.options.MaxDumpBodySize +} + func (t *HTTPServer) loglayer(handler http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fullRequest, _ := httputil.DumpRequest(r, true) - lrw := newLoggingResponseWriter(w) + var fullRequest []byte + if t.shouldDumpBody(r.ContentLength) { + fullRequest, _ = httputil.DumpRequest(r, false) + } else { + fullRequest, _ = httputil.DumpRequest(r, true) + } + lrw := newLoggingResponseWriter(w, t.options.MaxDumpBodySize) handler.ServeHTTP(lrw, r) // Handles file write if enabled @@ -52,7 +62,7 @@ func (t *HTTPServer) loglayer(handler http.Handler) http.Handler { err error ) if t.options.Sandbox { - maxFileSize := toMb(t.options.MaxFileSize) + maxFileSize := unit.ToMb(t.options.MaxFileSize) // check header content length if r.ContentLength > maxFileSize { gologger.Print().Msg("request too large") @@ -81,24 +91,29 @@ func (t *HTTPServer) loglayer(handler http.Handler) http.Handler { lrw.Header().Write(headers) //nolint gologger.Print().Msgf("\nRemote Address: %s\n%s\n%s %d %s\n%s\n%s\n", r.RemoteAddr, string(fullRequest), r.Proto, lrw.statusCode, http.StatusText(lrw.statusCode), headers.String(), string(lrw.Data)) } else { - gologger.Print().Msgf("%s \"%s %s %s\" %d %d", r.RemoteAddr, r.Method, r.URL, r.Proto, lrw.statusCode, len(lrw.Data)) + gologger.Print().Msgf("%s \"%s %s %s\" %d %d", r.RemoteAddr, r.Method, r.URL, r.Proto, lrw.statusCode, lrw.Size) } }) } type loggingResponseWriter struct { http.ResponseWriter - statusCode int - Data []byte + statusCode int + Data []byte + Size int + MaxDumpSize int64 } -func newLoggingResponseWriter(w http.ResponseWriter) *loggingResponseWriter { - return &loggingResponseWriter{w, http.StatusOK, []byte{}} +func newLoggingResponseWriter(w http.ResponseWriter, maxSize int64) *loggingResponseWriter { + return &loggingResponseWriter{w, http.StatusOK, []byte{}, 0, maxSize} } // Write the data func (lrw *loggingResponseWriter) Write(data []byte) (int, error) { - lrw.Data = append(lrw.Data, data...) + if len(lrw.Data) < int(lrw.MaxDumpSize) { + lrw.Data = append(lrw.Data, data...) + } + lrw.Size += len(data) return lrw.ResponseWriter.Write(data) } diff --git a/pkg/httpserver/util.go b/pkg/httpserver/util.go deleted file mode 100644 index 4c69d6f..0000000 --- a/pkg/httpserver/util.go +++ /dev/null @@ -1,5 +0,0 @@ -package httpserver - -func toMb(n int) int64 { - return int64(n) * 1024 * 1024 -} diff --git a/pkg/unit/unit.go b/pkg/unit/unit.go new file mode 100644 index 0000000..98cdb35 --- /dev/null +++ b/pkg/unit/unit.go @@ -0,0 +1,6 @@ +package unit + +// ToMb converts bytes to megabytes +func ToMb(n int) int64 { + return int64(n) * 1024 * 1024 +}