github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/protocols/headless/engine/hijack.go (about)

     1  package engine
     2  
     3  // TODO: redundant from katana - the whole headless package should be replace with katana
     4  
     5  import (
     6  	"encoding/base64"
     7  
     8  	"github.com/go-rod/rod"
     9  	"github.com/go-rod/rod/lib/proto"
    10  )
    11  
    12  // NewHijack create hijack from page.
    13  func NewHijack(page *rod.Page) *Hijack {
    14  	return &Hijack{
    15  		page:    page,
    16  		disable: &proto.FetchDisable{},
    17  	}
    18  }
    19  
    20  // HijackHandler type
    21  type HijackHandler = func(e *proto.FetchRequestPaused) error
    22  
    23  // Hijack is a hijack handler
    24  type Hijack struct {
    25  	page    *rod.Page
    26  	enable  *proto.FetchEnable
    27  	disable *proto.FetchDisable
    28  	cancel  func()
    29  }
    30  
    31  // SetPattern set pattern directly
    32  func (h *Hijack) SetPattern(pattern *proto.FetchRequestPattern) {
    33  	h.enable = &proto.FetchEnable{
    34  		Patterns: []*proto.FetchRequestPattern{pattern},
    35  	}
    36  }
    37  
    38  // Start hijack.
    39  func (h *Hijack) Start(handler HijackHandler) func() error {
    40  	if h.enable == nil {
    41  		panic("hijack pattern not set")
    42  	}
    43  
    44  	p, cancel := h.page.WithCancel()
    45  	h.cancel = cancel
    46  
    47  	err := h.enable.Call(p)
    48  	if err != nil {
    49  		return func() error { return err }
    50  	}
    51  
    52  	wait := p.EachEvent(func(e *proto.FetchRequestPaused) {
    53  		if handler != nil {
    54  			err = handler(e)
    55  		}
    56  	})
    57  
    58  	return func() error {
    59  		wait()
    60  		return err
    61  	}
    62  }
    63  
    64  // Stop
    65  func (h *Hijack) Stop() error {
    66  	if h.cancel != nil {
    67  		h.cancel()
    68  	}
    69  	return h.disable.Call(h.page)
    70  }
    71  
    72  // FetchGetResponseBody get request body.
    73  func FetchGetResponseBody(page *rod.Page, e *proto.FetchRequestPaused) ([]byte, error) {
    74  	m := proto.FetchGetResponseBody{
    75  		RequestID: e.RequestID,
    76  	}
    77  	r, err := m.Call(page)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  
    82  	if !r.Base64Encoded {
    83  		return []byte(r.Body), nil
    84  	}
    85  
    86  	bs, err := base64.StdEncoding.DecodeString(r.Body)
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  	return bs, nil
    91  }
    92  
    93  // FetchContinueRequest continue request
    94  func FetchContinueRequest(page *rod.Page, e *proto.FetchRequestPaused) error {
    95  	m := proto.FetchContinueRequest{
    96  		RequestID: e.RequestID,
    97  	}
    98  	return m.Call(page)
    99  }