mirror of
https://github.com/MontFerret/ferret.git
synced 2024-12-16 11:37:36 +02:00
34c8c02258
* Externalized HTML drivers * Fixed unit tests * Updated logging * Added support to set default driver * Updated GetIn and SetIn helpers
129 lines
2.7 KiB
Go
129 lines
2.7 KiB
Go
package http
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"net/http"
|
|
"net/url"
|
|
|
|
"github.com/MontFerret/ferret/pkg/drivers"
|
|
"github.com/MontFerret/ferret/pkg/drivers/common"
|
|
"github.com/MontFerret/ferret/pkg/runtime/logging"
|
|
"github.com/MontFerret/ferret/pkg/runtime/values"
|
|
"github.com/PuerkitoBio/goquery"
|
|
"github.com/corpix/uarand"
|
|
"github.com/pkg/errors"
|
|
"github.com/sethgrid/pester"
|
|
)
|
|
|
|
const DriverName = "http"
|
|
|
|
type Driver struct {
|
|
client *pester.Client
|
|
options *Options
|
|
}
|
|
|
|
func NewDriver(opts ...Option) *Driver {
|
|
drv := new(Driver)
|
|
drv.options = newOptions(opts)
|
|
|
|
if drv.options.proxy == "" {
|
|
drv.client = pester.New()
|
|
} else {
|
|
client, err := newClientWithProxy(drv.options)
|
|
|
|
if err != nil {
|
|
drv.client = pester.New()
|
|
} else {
|
|
drv.client = pester.NewExtendedClient(client)
|
|
}
|
|
}
|
|
|
|
drv.client.Concurrency = drv.options.concurrency
|
|
drv.client.MaxRetries = drv.options.maxRetries
|
|
drv.client.Backoff = drv.options.backoff
|
|
|
|
return drv
|
|
}
|
|
|
|
func newClientWithProxy(options *Options) (*http.Client, error) {
|
|
proxyURL, err := url.Parse(options.proxy)
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
proxy := http.ProxyURL(proxyURL)
|
|
tr := &http.Transport{Proxy: proxy}
|
|
|
|
return &http.Client{Transport: tr}, nil
|
|
}
|
|
|
|
func (drv *Driver) Name() string {
|
|
return DriverName
|
|
}
|
|
|
|
func (drv *Driver) GetDocument(ctx context.Context, targetURL values.String) (drivers.HTMLDocument, error) {
|
|
u := targetURL.String()
|
|
req, err := http.NewRequest(http.MethodGet, u, nil)
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8")
|
|
req.Header.Set("Accept-Language", "en-US,en;q=0.9,ru;q=0.8")
|
|
req.Header.Set("Cache-Control", "no-cache")
|
|
req.Header.Set("Pragma", "no-cache")
|
|
|
|
req = req.WithContext(ctx)
|
|
|
|
ua := common.GetUserAgent(drv.options.userAgent)
|
|
|
|
logger := logging.FromContext(ctx)
|
|
logger.
|
|
Debug().
|
|
Timestamp().
|
|
Str("user-agent", ua).
|
|
Msg("using User-Agent")
|
|
|
|
// use custom user agent
|
|
if ua != "" {
|
|
req.Header.Set("User-Agent", uarand.GetRandom())
|
|
}
|
|
|
|
resp, err := drv.client.Do(req)
|
|
|
|
if err != nil {
|
|
return nil, errors.Wrapf(err, "failed to retrieve a document %s", u)
|
|
}
|
|
|
|
defer resp.Body.Close()
|
|
|
|
doc, err := goquery.NewDocumentFromReader(resp.Body)
|
|
|
|
if err != nil {
|
|
return nil, errors.Wrapf(err, "failed to parse a document %s", u)
|
|
}
|
|
|
|
return NewHTMLDocument(u, doc)
|
|
}
|
|
|
|
func (drv *Driver) ParseDocument(_ context.Context, str values.String) (drivers.HTMLDocument, error) {
|
|
buf := bytes.NewBuffer([]byte(str))
|
|
|
|
doc, err := goquery.NewDocumentFromReader(buf)
|
|
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "failed to parse a document")
|
|
}
|
|
|
|
return NewHTMLDocument("#string", doc)
|
|
}
|
|
|
|
func (drv *Driver) Close() error {
|
|
drv.client = nil
|
|
|
|
return nil
|
|
}
|