Skip to content

Commit

Permalink
(feat) add immediate screenshot mode and skip screenshot flags
Browse files Browse the repository at this point in the history
  • Loading branch information
leonjza committed Sep 18, 2024
1 parent f4ca218 commit c4478e2
Show file tree
Hide file tree
Showing 12 changed files with 345 additions and 24 deletions.
3 changes: 2 additions & 1 deletion cmd/scan.go
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ flags.`)),
}

if len(scanWriters) == 0 {
log.Warn("no writers have been configured. only saving screenshots. add writers using --write-* flags")
log.Warn("no writers have been configured. to persist probe results, add writers using --write-* flags")
}

// Get the runner up. Basically, all of the subcommands will use this.
Expand Down Expand Up @@ -135,6 +135,7 @@ func init() {
scanCmd.PersistentFlags().StringVarP(&opts.Scan.ScreenshotPath, "screenshot-path", "s", "./screenshots", "Path to store screenshots")
scanCmd.PersistentFlags().StringVar(&opts.Scan.ScreenshotFormat, "screenshot-format", "jpeg", "Format to save screenshots as. Valid formats are: jpeg, png")
scanCmd.PersistentFlags().BoolVar(&opts.Scan.ScreenshotFullPage, "screenshot-fullpage", false, "Do full-page screenshots, instead of just the viewport")
scanCmd.PersistentFlags().BoolVar(&opts.Scan.ScreenshotSkipSave, "screenshot-skip-save", false, "Do not save screenshots to the screenshot-path (useful together with --write-screenshots)")
scanCmd.PersistentFlags().StringVar(&opts.Scan.JavaScript, "javascript", "", "A JavaScript function to evaluate on every page, before a screenshot. Note: It must be a JavaScript function! e.g., () => console.log('gowitness');")
scanCmd.PersistentFlags().StringVar(&opts.Scan.JavaScriptFile, "javascript-file", "", "A file containing a JavaScript function to evaluate on every page, before a screenshot. See --javascript")
scanCmd.PersistentFlags().BoolVar(&opts.Scan.SaveContent, "save-content", false, "Save content from network requests to the configured writers. WARNING: This flag has the potential to make your storage explode in size")
Expand Down
18 changes: 10 additions & 8 deletions pkg/runner/drivers/chromedp.go
Original file line number Diff line number Diff line change
Expand Up @@ -430,14 +430,16 @@ func (run *Chromedp) Witness(target string, runner *runner.Runner) (*models.Resu
result.Screenshot = base64.StdEncoding.EncodeToString(img)
}

// write the screenshot to disk
result.Filename = islazy.SafeFileName(target) + "." + run.options.Scan.ScreenshotFormat
result.Filename = islazy.LeftTrucate(result.Filename, 200)
if err := os.WriteFile(
filepath.Join(run.options.Scan.ScreenshotPath, result.Filename),
img, os.FileMode(0664),
); err != nil {
return nil, fmt.Errorf("could not write screenshot to disk: %w", err)
// write the screenshot to disk if we have a path
if !run.options.Scan.ScreenshotSkipSave {
result.Filename = islazy.SafeFileName(target) + "." + run.options.Scan.ScreenshotFormat
result.Filename = islazy.LeftTrucate(result.Filename, 200)
if err := os.WriteFile(
filepath.Join(run.options.Scan.ScreenshotPath, result.Filename),
img, os.FileMode(0664),
); err != nil {
return nil, fmt.Errorf("could not write screenshot to disk: %w", err)
}
}

// calculate and set the perception hash
Expand Down
18 changes: 10 additions & 8 deletions pkg/runner/drivers/go-rod.go
Original file line number Diff line number Diff line change
Expand Up @@ -430,14 +430,16 @@ func (run *Gorod) Witness(target string, runner *runner.Runner) (*models.Result,
result.Screenshot = base64.StdEncoding.EncodeToString(img)
}

// write the screenshot to disk
result.Filename = islazy.SafeFileName(target) + "." + run.options.Scan.ScreenshotFormat
result.Filename = islazy.LeftTrucate(result.Filename, 200)
if err := os.WriteFile(
filepath.Join(run.options.Scan.ScreenshotPath, result.Filename),
img, os.FileMode(0664),
); err != nil {
return nil, fmt.Errorf("could not write screenshot to disk: %w", err)
// write the screenshot to disk if we have a path
if !run.options.Scan.ScreenshotSkipSave {
result.Filename = islazy.SafeFileName(target) + "." + run.options.Scan.ScreenshotFormat
result.Filename = islazy.LeftTrucate(result.Filename, 200)
if err := os.WriteFile(
filepath.Join(run.options.Scan.ScreenshotPath, result.Filename),
img, os.FileMode(0664),
); err != nil {
return nil, fmt.Errorf("could not write screenshot to disk: %w", err)
}
}

// calculate and set the perception hash
Expand Down
6 changes: 5 additions & 1 deletion pkg/runner/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -70,14 +70,18 @@ type Scan struct {
UriFilter []string
// Don't write HTML response content
SkipHTML bool
// ScreenshotPath is the path where screenshot images will be stored
// ScreenshotPath is the path where screenshot images will be stored.
// An empty value means drivers will not write screenshots to disk. In
// that case, you'd need to specify writer saves.
ScreenshotPath string
// ScreenshotFormat to save as
ScreenshotFormat string
// ScreenshotFullPage saves full, scrolled web pages
ScreenshotFullPage bool
// ScreenshotToWriter passes screenshots as a model property to writers
ScreenshotToWriter bool
// ScreenshotSkipSave skips saving screenshots to disk
ScreenshotSkipSave bool
// JavaScript to evaluate on every page
JavaScript string
JavaScriptFile string
Expand Down
14 changes: 9 additions & 5 deletions pkg/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,16 @@ type Runner struct {
// New gets a new Runner ready for probing.
// It's up to the caller to call Close() on the runner
func NewRunner(logger *slog.Logger, driver Driver, opts Options, writers []writers.Writer) (*Runner, error) {
screenshotPath, err := islazy.CreateDir(opts.Scan.ScreenshotPath)
if err != nil {
return nil, err
if !opts.Scan.ScreenshotSkipSave {
screenshotPath, err := islazy.CreateDir(opts.Scan.ScreenshotPath)
if err != nil {
return nil, err
}
opts.Scan.ScreenshotPath = screenshotPath
logger.Debug("final screenshot path", "screenshot-path", opts.Scan.ScreenshotPath)
} else {
logger.Debug("not saving screenshots to disk")
}
opts.Scan.ScreenshotPath = screenshotPath
logger.Debug("final screenshot path", "screenshot-path", opts.Scan.ScreenshotPath)

// screenshot format check
if !islazy.SliceHasStr([]string{"jpeg", "png"}, opts.Scan.ScreenshotFormat) {
Expand Down
78 changes: 78 additions & 0 deletions pkg/writers/memory.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
package writers

import (
"errors"
"sync"

"github.com/sensepost/gowitness/pkg/models"
)

// MemoryWriter is a memory-based results queue with a maximum slot count
type MemoryWriter struct {
slots int
results []*models.Result
mutex sync.Mutex
}

// NewMemoryWriter initializes a MemoryWriter with the specified number of slots
func NewMemoryWriter(slots int) (*MemoryWriter, error) {
if slots <= 0 {
return nil, errors.New("slots need to be a positive integer")
}

return &MemoryWriter{
slots: slots,
results: make([]*models.Result, 0, slots),
mutex: sync.Mutex{},
}, nil
}

// Write adds a new result to the MemoryWriter.
func (s *MemoryWriter) Write(result *models.Result) error {
s.mutex.Lock()
defer s.mutex.Unlock()

if len(s.results) >= s.slots {
s.results = s.results[1:]
}

s.results = append(s.results, result)

return nil
}

// GetLatest retrieves the most recently added result.
func (s *MemoryWriter) GetLatest() *models.Result {
s.mutex.Lock()
defer s.mutex.Unlock()

if len(s.results) == 0 {
return nil
}

return s.results[len(s.results)-1]
}

// GetFirst retrieves the oldest result in the MemoryWriter.
func (s *MemoryWriter) GetFirst() *models.Result {
s.mutex.Lock()
defer s.mutex.Unlock()

if len(s.results) == 0 {
return nil
}

return s.results[0]
}

// GetAllResults returns a copy of all current results.
func (s *MemoryWriter) GetAllResults() []*models.Result {
s.mutex.Lock()
defer s.mutex.Unlock()

// Create a copy to prevent external modification
resultsCopy := make([]*models.Result, len(s.results))
copy(resultsCopy, s.results)

return resultsCopy
}
2 changes: 1 addition & 1 deletion web/api/submit.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ func (h *ApiHandler) SubmitHandler(w http.ResponseWriter, r *http.Request) {
w.Write(jsonData)
}

// dispatchRunner run's a runner
// dispatchRunner run's a runner in a separate goroutine
func dispatchRunner(runner *runner.Runner, targets []string) {
// feed in targets
go func() {
Expand Down
101 changes: 101 additions & 0 deletions web/api/submit_single.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
package api

import (
"encoding/json"
"log/slog"
"net/http"

"github.com/sensepost/gowitness/pkg/log"
"github.com/sensepost/gowitness/pkg/runner"
driver "github.com/sensepost/gowitness/pkg/runner/drivers"
"github.com/sensepost/gowitness/pkg/writers"
)

type submitSingleRequest struct {
URL string `json:"url"`
Options *submitRequestOptions `json:"options"`
}

// SubmitSingleHandler submits a URL to scan, returning the result.
//
// @Summary Submit a single URL for probing
// @Description Starts a new probing routine for a URL and options, returning the results when done.
// @Tags Results
// @Accept json
// @Produce json
// @Param query body submitSingleRequest true "The URL scanning request object"
// @Success 200 {string} string "Probing started"
// @Router /submit/single [post]
func (h *ApiHandler) SubmitSingleHandler(w http.ResponseWriter, r *http.Request) {
var request submitSingleRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
log.Error("failed to read json request", "err", err)
http.Error(w, "Error reading JSON request", http.StatusInternalServerError)
return
}

if request.URL == "" {
http.Error(w, "No URL provided", http.StatusBadRequest)
return
}

options := runner.NewDefaultOptions()
options.Scan.ScreenshotToWriter = true
options.Scan.ScreenshotSkipSave = true

// Override default values with request options
if request.Options != nil {
if request.Options.X != 0 {
options.Chrome.WindowX = request.Options.X
}
if request.Options.Y != 0 {
options.Chrome.WindowY = request.Options.Y
}
if request.Options.UserAgent != "" {
options.Chrome.UserAgent = request.Options.UserAgent
}
if request.Options.Timeout != 0 {
options.Scan.Timeout = request.Options.Timeout
}
if request.Options.Format != "" {
options.Scan.ScreenshotFormat = request.Options.Format
}
}

writer, err := writers.NewMemoryWriter(1)
if err != nil {
http.Error(w, "Error getting a memory writer", http.StatusInternalServerError)
return
}

logger := slog.New(log.Logger)

driver, err := driver.NewChromedp(logger, *options)
if err != nil {
http.Error(w, "Error sarting driver", http.StatusInternalServerError)
return
}

runner, err := runner.NewRunner(logger, driver, *options, []writers.Writer{writer})
if err != nil {
log.Error("error starting runner", "err", err)
http.Error(w, "Error starting runner", http.StatusInternalServerError)
return
}

go func() {
runner.Targets <- request.URL
close(runner.Targets)
}()

runner.Run()
runner.Close()

jsonData, err := json.Marshal(writer.GetLatest())
if err != nil {
http.Error(w, "Error creating JSON response", http.StatusInternalServerError)
return
}

w.Write(jsonData)
}
48 changes: 48 additions & 0 deletions web/docs/docs.go
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,40 @@ const docTemplate = `{
}
}
},
"/submit/single": {
"post": {
"description": "Starts a new probing routine for a URL and options, returning the results when done.",
"consumes": [
"application/json"
],
"produces": [
"application/json"
],
"tags": [
"Results"
],
"summary": "Submit a single URL for probing",
"parameters": [
{
"description": "The URL scanning request object",
"name": "query",
"in": "body",
"required": true,
"schema": {
"$ref": "#/definitions/api.submitSingleRequest"
}
}
],
"responses": {
"200": {
"description": "Probing started",
"schema": {
"type": "string"
}
}
}
}
},
"/wappalyzer": {
"get": {
"description": "Get all of the available wappalyzer data.",
Expand Down Expand Up @@ -549,6 +583,17 @@ const docTemplate = `{
}
}
},
"api.submitSingleRequest": {
"type": "object",
"properties": {
"options": {
"$ref": "#/definitions/api.submitRequestOptions"
},
"url": {
"type": "string"
}
}
},
"api.technologyListResponse": {
"type": "object",
"properties": {
Expand Down Expand Up @@ -746,6 +791,9 @@ const docTemplate = `{
"perception_hash": {
"type": "string"
},
"perception_hash_group_id": {
"type": "integer"
},
"probed_at": {
"type": "string"
},
Expand Down
Loading

0 comments on commit c4478e2

Please sign in to comment.