Add/Expose Headless action that waits until an element appears #1096

* add timeout and pollTime options

v2/pkg/protocols/headless/engine/page_actions.go

@@ -1,6 +1,7 @@
 package engine
 
 import (
+	"context"
 	"io/ioutil"
 	"net"
 	"net/url"
@@ -11,6 +12,7 @@ import (
 
 	"github.com/go-rod/rod"
 	"github.com/go-rod/rod/lib/proto"
+	"github.com/go-rod/rod/lib/utils"
 	"github.com/pkg/errors"
 	"github.com/segmentio/ksuid"
 	"github.com/valyala/fasttemplate"
@@ -46,7 +48,7 @@ func (p *Page) ExecuteActions(baseURL *url.URL, actions []*Action) (map[string]s
 		case ActionGetResource:
 			err = p.GetResource(act, outData)
 		case ActionExtract:
-			err = p.SelectInputElement(act, outData)
+			err = p.ExtractElement(act, outData)
 		case ActionWaitEvent:
 			err = p.WaitEvent(act, outData)
 		case ActionFilesInput:
@@ -85,18 +87,83 @@ type requestRule struct {
 	Args   map[string]string
 }
 
+const elementDidNotAppearMessage = "Element did not appear in the given amount of time"
+
 // WaitVisible waits until an element appears.
 func (p *Page) WaitVisible(act *Action, out map[string]string) error {
-	element, err := p.pageElementBy(act.Data)
+	timeout, err := getTimeout(act)
 	if err != nil {
-		return errors.Wrap(err, "could not find element")
+		return errors.Wrap(err, "Wrong timeout given")
 	}
-	if err = element.WaitVisible(); err != nil {
+
-		return errors.Wrap(err, "could not wait element")
+	pollTime, err := getPollTime(act)
+	if err != nil {
+		return errors.Wrap(err, "Wrong polling time given")
 	}
+
+	element, _ := p.Sleeper(pollTime, timeout).
+		Timeout(timeout).
+		pageElementBy(act.Data)
+
+	if element != nil {
+		if err := element.WaitVisible(); err != nil {
+			return errors.Wrap(err, elementDidNotAppearMessage)
+		}
+	} else {
+		return errors.New(elementDidNotAppearMessage)
+	}
+
 	return nil
 }
 
+func (p *Page) Sleeper(pollTimeout, timeout time.Duration) *Page {
+	page := *p
+	page.page = page.Page().Sleeper(func() utils.Sleeper {
+		return createBackOffSleeper(pollTimeout, timeout)
+	})
+	return &page
+}
+
+func (p *Page) Timeout(timeout time.Duration) *Page {
+	page := *p
+	page.page = page.Page().Timeout(timeout)
+	return &page
+}
+
+func createBackOffSleeper(pollTimeout, timeout time.Duration) utils.Sleeper {
+	backoffSleeper := utils.BackoffSleeper(pollTimeout, timeout, func(duration time.Duration) time.Duration {
+		return duration
+	})
+
+	return func(ctx context.Context) error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		return backoffSleeper(ctx)
+	}
+}
+
+func getTimeout(act *Action) (time.Duration, error) {
+	return geTimeParameter(act, "timeout", 3, time.Second)
+}
+
+func getPollTime(act *Action) (time.Duration, error) {
+	return geTimeParameter(act, "pollTime", 100, time.Millisecond)
+}
+
+func geTimeParameter(act *Action, parameterName string, defaultValue time.Duration, duration time.Duration) (time.Duration, error) {
+	pollTimeString := act.GetArg(parameterName)
+	if pollTimeString == "" {
+		return defaultValue * duration, nil
+	}
+	timeout, err := strconv.Atoi(pollTimeString)
+	if err != nil {
+		return time.Duration(0), err
+	}
+	return time.Duration(timeout) * duration, nil
+}
+
 // ActionAddHeader executes a AddHeader action.
 func (p *Page) ActionAddHeader(act *Action, out map[string]string /*TODO review unused parameter*/) error {
 	in := act.GetArg("part")

v2/pkg/protocols/headless/engine/page_actions_test.go

@@ -9,9 +9,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
-	"github.com/stretchr/testify/require"
 )
 
 func TestActionNavigate(t *testing.T) {
@@ -281,17 +282,24 @@ func TestActionHeadersChange(t *testing.T) {
 
 func TestActionWaitVisible(t *testing.T) {
 	t.Run("wait for an element being visible", func(t *testing.T) {
-		_ = protocolstate.Init(&types.Options{})
+		testWaitVisible(t, 2*time.Second, func(page *Page, err error) {
+			require.Nil(t, err, "could not run page actions")
 
-		browser, err := New(&types.Options{ShowBrowser: false})
+			page.Page().MustElement("button").MustVisible()
-		require.Nil(t, err, "could not create browser")
+			page.Close()
-		defer browser.Close()
+		})
+	})
 
-		instance, err := browser.NewInstance()
+	t.Run("timeout because of element not visible", func(t *testing.T) {
-		require.Nil(t, err, "could not create browser instance")
+		testWaitVisible(t, time.Second/2, func(page *Page, err error) {
+			require.Error(t, err)
+			require.Contains(t, err.Error(), "Element did not appear in the given amount of time")
+		})
+	})
+}
 
-		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+func testWaitVisible(t *testing.T, timeout time.Duration, assert func(page *Page, err error)) {
-			fmt.Fprintln(w, `
+	response := `
 		<html>
 			<head>
 				<title>Nuclei Test Page</title>
@@ -300,54 +308,29 @@ func TestActionWaitVisible(t *testing.T) {
 			<script>
 				setTimeout(() => document.querySelector('#test').style.display = '', 1000);
 			</script>
-		</html>`)
+		</html>`
-		}))
-		defer ts.Close()
 
-		parsed, err := url.Parse(ts.URL)
+	_ = protocolstate.Init(&types.Options{})
-		require.Nil(t, err, "could not parse URL")
 
-		actions := []*Action{
+	browser, err := New(&types.Options{ShowBrowser: false})
-			{ActionType: "navigate", Data: map[string]string{"url": "{{BaseURL}}"}},
+	require.Nil(t, err, "could not create browser")
-			{ActionType: "waitvisible", Data: map[string]string{"by": "x", "xpath": "//button[@id='test']"}},
+	defer browser.Close()
-		}
-		_, page, err := instance.Run(parsed, actions, 20*time.Second)
-		require.Nil(t, err, "could not run page actions")
-		defer page.Close()
 
-		page.Page().MustElement("button").MustVisible()
+	instance, err := browser.NewInstance()
-	})
+	require.Nil(t, err, "could not create browser instance")
 
-	t.Run("timeout because of element not visible", func(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		_ = protocolstate.Init(&types.Options{})
+		fmt.Fprintln(w, response)
+	}))
+	defer ts.Close()
 
-		browser, err := New(&types.Options{ShowBrowser: false})
+	parsed, err := url.Parse(ts.URL)
-		require.Nil(t, err, "could not create browser")
+	require.Nil(t, err, "could not parse URL")
-		defer browser.Close()
 
-		instance, err := browser.NewInstance()
+	actions := []*Action{
-		require.Nil(t, err, "could not create browser instance")
+		{ActionType: "navigate", Data: map[string]string{"url": "{{BaseURL}}"}},
-
+		{ActionType: "waitvisible", Data: map[string]string{"by": "x", "xpath": "//button[@id='test']"}},
-		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	}
-			fmt.Fprintln(w, `
+	_, page, err := instance.Run(parsed, actions, timeout)
-		<html>
+	assert(page, err)
-			<head>
-				<title>Nuclei Test Page</title>
-			</head>
-			<button style="display:none" id="test">Wait for me!</button>
-		</html>`)
-		}))
-		defer ts.Close()
-
-		parsed, err := url.Parse(ts.URL)
-		require.Nil(t, err, "could not parse URL")
-
-		actions := []*Action{
-			{ActionType: "navigate", Data: map[string]string{"url": "{{BaseURL}}"}},
-			{ActionType: "waitvisible", Data: map[string]string{"by": "x", "xpath": "//button[@id='test']"}},
-		}
-		_, _, err = instance.Run(parsed, actions, 2*time.Second)
-		require.Error(t, err)
-		require.Contains(t, err.Error(), "could not wait element")
-	})
 }