1package plugins
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"io"
  7	"net/http"
  8	"net/url"
  9	"strings"
 10	"time"
 11
 12	"golang.org/x/net/html"
 13	"gopkg.in/irc.v3"
 14)
 15
 16func init() {
 17	Register(Search{})
 18}
 19
 20type Search struct{}
 21
 22func (Search) Triggers() []string {
 23	return []string{".ddg", ".g", ".mdn", ".wiki"}
 24}
 25
 26var ddgClient = &http.Client{
 27	Timeout: 10 * time.Second,
 28}
 29
 30func ddgParse(body io.Reader) (string, error) {
 31	inResult := false
 32	href := ""
 33	var snippet strings.Builder
 34	tok := html.NewTokenizer(body)
 35
 36	for tt := tok.Next(); tt != html.ErrorToken; tt = tok.Next() {
 37		switch tt {
 38		case html.StartTagToken:
 39			tag, hasAttr := tok.TagName()
 40
 41			if tag[0] == 'a' {
 42				if hasAttr {
 43					var key []byte
 44					var val []byte
 45					var lhref []byte
 46					isRes := false
 47
 48					for hasAttr {
 49						key, val, hasAttr = tok.TagAttr()
 50						if bytes.Equal(key, []byte("href")) {
 51							lhref = val
 52						} else if bytes.Equal(val, []byte("result__snippet")) {
 53							isRes = true
 54						}
 55					}
 56
 57					if isRes && len(lhref) != 0 {
 58						realUrl, err := url.Parse(string(lhref))
 59						if err == nil {
 60							if v, ok := realUrl.Query()["uddg"]; len(v) > 0 && ok {
 61								inResult = true
 62								href = v[0]
 63							}
 64						}
 65					}
 66				}
 67			} else if tag[0] == 'b' { // support bold text
 68				//snippet.WriteRune('\x02')
 69			}
 70		case html.TextToken:
 71			if inResult && snippet.Len() < 300 {
 72				snippet.Write(tok.Text())
 73			}
 74		case html.EndTagToken:
 75			tag, _ := tok.TagName()
 76
 77			if inResult && tag[0] == 'a' {
 78				var len int
 79				if snippet.Len() > 300 {
 80					len = 300
 81				} else {
 82					len = snippet.Len()
 83				}
 84				return snippet.String()[:len] + " - " + href, nil
 85			} else if inResult && tag[0] == 'b' {
 86				//snippet.WriteRune('\x02')
 87			}
 88		}
 89	}
 90
 91	if err := tok.Err(); err != io.EOF && err != nil {
 92		return "HTML parse error", err
 93	} else {
 94		return "No results.", nil
 95	}
 96}
 97
 98// Use DuckDuckGo's
 99func ddg(query string) (string, error) {
100	req, err := http.NewRequest("GET", "https://html.duckduckgo.com/html", nil)
101
102	if err != nil {
103		return "Client request error", err
104	}
105
106	req.Header.Add("User-Agent", "github.com/icyphox/Taigobot")
107
108	q := req.URL.Query()
109	q.Add("q", query)
110	req.URL.RawQuery = q.Encode()
111
112	res, err := ddgClient.Do(req)
113	if err != nil {
114		return "Server response error", err
115	}
116
117	defer res.Body.Close()
118	result, err := ddgParse(res.Body)
119	if err != nil {
120		return "HTML parse error", err
121	}
122	return result, nil
123}
124
125// This is just an alias for now.
126func google(query string) (string, error) {
127	return ddg(query)
128}
129
130func mdn(query string) (string, error) {
131	return ddg("site:https://developer.mozilla.org/en-US " + query)
132}
133
134func wiki(query string) (string, error) {
135	return ddg("site:https://en.wikipedia.org " + query)
136}
137
138func (Search) Execute(m *irc.Message) (string, error) {
139	parsed := strings.SplitN(m.Trailing(), " ", 2)
140	if len(parsed) != 2 {
141		return fmt.Sprintf("Usage: %s <query>", parsed[0]), nil
142	}
143	trigger, query := parsed[0], parsed[1]
144
145	switch trigger {
146	case ".mdn":
147		return mdn(query)
148	case ".wiki":
149		return wiki(query)
150	case ".g":
151		return google(query)
152	default:
153		return ddg(query)
154	}
155}