all repos — paprika @ ed37b272f5791d139f3031bd35c8b3c41592f356

go rewrite of taigabot

plugins/quotes.go (view raw)

  1package plugins
  2
  3import (
  4	"bytes"
  5	"errors"
  6	"fmt"
  7	"log"
  8	"math/rand"
  9	"strconv"
 10	"strings"
 11
 12	"git.icyphox.sh/paprika/database"
 13	"github.com/dgraph-io/badger/v3"
 14	"gopkg.in/irc.v3"
 15)
 16
 17type Quotes struct{}
 18
 19func init() {
 20	Register(Quotes{})
 21}
 22
 23func (Quotes) Triggers() []string {
 24	return []string{".q", ".quote"}
 25}
 26
 27func result(quoteNum int, total int, nick string, quote string) string {
 28	return fmt.Sprintf("[%d/%d] <%s> %s", quoteNum, total, nick, quote)
 29}
 30
 31var found = errors.New("Found")
 32var KeyEncodingError = errors.New("Unexpected key encoding")
 33var TooManyQuotes = errors.New("Too many quotes")
 34
 35func getQuoteTotal(txn *badger.Txn, keyPrefix []byte) (int, error) {
 36	item, err := txn.Get([]byte(keyPrefix))
 37	if err != nil {
 38		return 0, err
 39	}
 40	it, err := item.ValueCopy(nil)
 41	if err != nil {
 42		return 0, err
 43	}
 44	res, err := strconv.Atoi(string(it))
 45	if _, ok := err.(*strconv.NumError); ok {
 46		log.Printf("quotes.go: Warning: Something is wrong with the value in key: %s", keyPrefix)
 47		return 0, nil // return 0 in hopes of it being overwritten
 48	} else if err != nil {
 49		return 0, err
 50	}
 51	return res, nil
 52}
 53
 54func getAndIncrementQuoteTotal(txn *badger.Txn, keyPrefix []byte) (int, error) {
 55	total, err := getQuoteTotal(txn, keyPrefix)
 56	if err == badger.ErrKeyNotFound {
 57		total = 0
 58	} else if err != nil {
 59		return 0, err
 60	}
 61	total++
 62	return total, txn.Set(keyPrefix, []byte(strconv.Itoa(total)))
 63}
 64
 65func findQuotes(nick string, keyPrefix, search []byte) (string, error) {
 66	var (
 67		num   int
 68		total int
 69		quote string
 70	)
 71
 72	err := database.DB.DB.View(func(txn *badger.Txn) error {
 73		iter := txn.NewIterator(badger.DefaultIteratorOptions)
 74		defer iter.Close()
 75
 76		var err error
 77		prefix := append(keyPrefix, ' ')
 78		for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
 79			item := iter.Item()
 80			key := item.Key()
 81			err = item.Value(func(val []byte) error {
 82				if bytes.Contains(val, search) {
 83					quote = string(val)
 84					keys := bytes.SplitN(key, []byte{' '}, 3)
 85					if len(keys) != 3 {
 86						log.Printf("quotes.go: Key Error: %s is not in expected format", key)
 87						return KeyEncodingError
 88					}
 89					num, err = database.DecodeNumber(keys[2])
 90					if err != nil {
 91						return err
 92					} else {
 93						return found
 94					}
 95				} else {
 96					return nil
 97				}
 98			})
 99
100			if err != nil {
101				break
102			}
103		}
104
105		if err == found {
106			total, err = getQuoteTotal(txn, []byte(keyPrefix))
107			if err != nil {
108				return err
109			} else {
110				return found
111			}
112		} else {
113			return err
114		}
115	})
116
117	if err == nil {
118		return "No quote found.", nil
119	} else if err == found {
120		return result(num, total, nick, quote), nil
121	} else {
122		return "", err
123	}
124}
125
126func addQuote(keyPrefix, quote []byte) (string, error) {
127	var id int
128	err := database.DB.DB.Update(func(txn *badger.Txn) error {
129		var err error
130		id, err = getAndIncrementQuoteTotal(txn, keyPrefix)
131		if err != nil {
132			return err
133		} else if id > 5_000 {
134			return TooManyQuotes
135		}
136
137		encodedId, err := database.EncodeNumber(id)
138		if err != nil {
139			return err
140		}
141
142		key := append(keyPrefix, ' ')
143		key = append(key, encodedId...)
144
145		err = txn.Set(key, quote)
146		if err != nil {
147			return err
148		}
149		return nil
150	})
151	if err == nil {
152		return fmt.Sprintf("Quote %d added.", id), err
153	} else {
154		return "", err
155	}
156}
157
158func getQuote(nick string, qnum int, keyPrefix []byte) (string, error) {
159	var (
160		num      int
161		total    int
162		quote    string
163		negative bool
164	)
165
166	if qnum < 0 {
167		qnum += 1
168		negative = true
169	} else {
170		negative = false
171	}
172
173	err := database.DB.DB.View(func(txn *badger.Txn) error {
174		var err error
175		total, err = getQuoteTotal(txn, keyPrefix)
176		if err != nil {
177			return err
178		}
179
180		num = qnum
181		if negative {
182			num = total + qnum
183			if num < 1 {
184				return badger.ErrKeyNotFound
185			}
186		} else if num > total {
187			return badger.ErrKeyNotFound
188		} else if num == randomQuote {
189			// [1, total+1)
190			num = rand.Intn(total) + 1
191		}
192
193		encodeQnum, err := database.EncodeNumber(num)
194		if err != nil {
195			return err
196		}
197		encodedKey := append(keyPrefix, ' ')
198		encodedKey = append(encodedKey, encodeQnum...)
199
200		qItem, err := txn.Get(encodedKey)
201		if err != nil {
202			return err
203		}
204		quoteT, err := qItem.ValueCopy(nil)
205		if err != nil {
206			return err
207		}
208		quote = string(quoteT)
209		return nil
210	})
211
212	if err == badger.ErrKeyNotFound {
213		return "No such quote for " + nick, nil
214	} else if err != nil {
215		return "", err
216	} else {
217		return result(num, total, nick, quote), nil
218	}
219}
220
221const randomQuote = 0
222const (
223	addQ int = iota
224	getQ
225	start
226	parseNick
227	parseGetParam
228)
229
230func (Quotes) Execute(m *irc.Message) (string, error) {
231	params := strings.Split(m.Trailing(), " ")
232	if len(params) == 1 {
233		return ".q [ add ] nickname [ quote | search | number ]", nil
234	}
235
236	pState := start
237	cState := getQ
238
239	var nick string
240	keyPrefix := []byte(m.Params[0] + " ")
241	for i := 1; i < len(params); i++ {
242		word := params[i]
243	back:
244		if len(word) == 0 {
245			continue
246		}
247		switch pState {
248		case start:
249			pState = parseNick
250			if word == "add" {
251				cState = addQ
252			} else {
253				goto back
254			}
255		case parseNick:
256			if word == "<" || len(word) == 0 {
257				break
258			}
259
260			// <xyz> -> xyz
261			word = strings.TrimPrefix(word, "<")
262			word = strings.TrimSuffix(word, ">")
263			if len(word) == 0 {
264				return "Invalid nickname given", nil
265			}
266			// This is used elsewhere to check the prefix of a "target"
267			// if it's true, then this word still has a prefix we can
268			// remove.
269			if likelyInvalidNickChr(word[0]) {
270				word = word[1:]
271			}
272			if len(word) == 0 {
273				return "Invalid nickname given", nil
274			}
275			// we only allow "< " prefix, not "<" + 2*sym
276			for j := 0; j < len(word); j++ {
277				if likelyInvalidNickChr(word[j]) {
278					return fmt.Sprintf("Invalid nickname: %s", word), nil
279				}
280			}
281			nick = word
282			keyPrefix = append(keyPrefix, nick...)
283			if cState == addQ {
284				quote := strings.Join(params[i+1:], " ")
285				if len(quote) == 0 {
286					return "Empty quote not allowed.", nil
287				}
288				return addQuote(keyPrefix, []byte(quote))
289			} else {
290				pState = parseGetParam
291			}
292		case parseGetParam:
293			if i+1 == len(params) {
294				qnum, err := strconv.Atoi(word)
295				if err != nil {
296					return findQuotes(nick, keyPrefix, []byte(word))
297				} else {
298					return getQuote(nick, qnum, keyPrefix)
299				}
300			} else {
301				quote := strings.Join(params[i+1:], " ")
302				return findQuotes(nick, keyPrefix, []byte(quote))
303			}
304		}
305	}
306	// If no number given, use 0 to indicate random quote.
307	if pState == parseGetParam {
308		return getQuote(nick, randomQuote, keyPrefix)
309	} else {
310		return "Invalid number of parameters.", nil
311	}
312}