scope/search/web/search.go
2021-12-08 09:24:05 -08:00

152 lines
2.7 KiB
Go

package web
import (
"net/http"
"sort"
"sync"
"time"
"golang.org/x/sync/errgroup"
)
func init() {
http.DefaultClient.Timeout = 5 * time.Second
}
type Result struct {
Title string
Link string
Desc string
Engines []string
Rank int
}
type Engine interface {
// Set search keyword for engine
SetKeyword(string)
// Set User Agent. If string is empty,
// an acceptable will should be used.
SetUserAgent(string)
// Set page number to search
SetPage(int)
// Initialize engine (make requests, set variables, etc.)
Init() error
// Run function for each search result,
// inputting index
Each(func(int) error) error
// Get title from index given by Each()
Title(int) (string, error)
// Get link from index given by Each()
Link(int) (string, error)
// Get description from index given by Each()
Desc(int) (string, error)
// Return shortened name of search engine.
// Should be lowercase (e.g. google, ddg, bing)
Name() string
}
type Options struct {
Keyword string
UserAgent string
Page int
}
func Search(opts Options, engines ...Engine) ([]*Result, error) {
var outMtx sync.Mutex
var out []*Result
wg := errgroup.Group{}
for index, engine := range engines {
curIndex, curEngine := index, engine
wg.Go(func() error {
curEngine.SetKeyword(opts.Keyword)
curEngine.SetUserAgent(opts.UserAgent)
curEngine.SetPage(opts.Page)
if err := curEngine.Init(); err != nil {
return err
}
err := curEngine.Each(func(i int) error {
link, err := curEngine.Link(i)
if err != nil {
return err
}
rank := (curIndex * 100) + i
index, exists := linkExists(out, link)
if exists {
out[index].Engines = append(out[index].Engines, curEngine.Name())
if rank < out[index].Rank {
out[index].Rank = rank
}
return nil
}
title, err := curEngine.Title(i)
if err != nil {
return err
}
desc, err := curEngine.Desc(i)
if err != nil {
return err
}
if title == "" || link == "" || desc == "" {
return nil
}
if len(desc) > 500 {
desc = desc[:500] + "..."
}
result := &Result{
Title: title,
Link: link,
Desc: desc,
Rank: rank,
}
result.Engines = append(result.Engines, curEngine.Name())
outMtx.Lock()
out = append(out, result)
outMtx.Unlock()
return nil
})
if err != nil {
return err
}
sort.Slice(out, func(i, j int) bool {
return out[i].Rank < out[j].Rank
})
return nil
})
}
if err := wg.Wait(); err != nil {
return out, err
}
return out, nil
}
func linkExists(results []*Result, link string) (int, bool) {
for index, result := range results {
if result.Link == link {
return index, true
}
}
return -1, false
}