CTF2021/irc_bot
Archived
4
0

Add missing db.go

This commit is contained in:
Denis Zheleztsov 2021-02-12 19:33:21 +03:00
parent 3e691011d4
commit 51c55ef708
Signed by: Difrex
GPG Key ID: D423378A747E202E

112
db.go Normal file
View File

@ -0,0 +1,112 @@
package main
import (
"bytes"
"encoding/binary"
"encoding/json"
"time"
"github.com/boltdb/bolt"
)
type LogLine struct {
Nick string
Message string
Time time.Time
ID uint64
}
type history struct {
db *bolt.DB
}
// Open creates a new db or open existing
func Open(path string) (*history, error) {
db, err := bolt.Open(path, 0600, nil)
if err != nil {
return nil, err
}
err = db.Update(func(tx *bolt.Tx) error {
if _, err := tx.CreateBucketIfNotExists([]byte("history")); err != nil {
return err
}
return nil
})
return &history{db}, err
}
func (h *history) WriteLog(l *LogLine) error {
err := h.db.Update(func(tx *bolt.Tx) error {
b, err := tx.CreateBucketIfNotExists([]byte("history"))
if err != nil {
return err
}
id, err := b.NextSequence()
if err != nil {
return err
}
l.ID = id
data, err := json.Marshal(l)
if err != nil {
return err
}
b.Put([]byte(l.Time.Format("2006-01-02T15:04:05")), data)
return nil
})
return err
}
func (h *history) GetAllLog() ([]LogLine, error) {
var logs []LogLine
err := h.db.View(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte("history"))
c := b.Cursor()
for k, v := c.First(); k != nil; k, v = c.Next() {
var l LogLine
if err := json.Unmarshal(v, &l); err != nil {
return err
}
logs = append(logs, l)
}
return nil
})
return logs, err
}
func (h *history) GetRangedLogs(from, to string) ([]LogLine, error) {
var logs []LogLine
err := h.db.View(func(tx *bolt.Tx) error {
c := tx.Bucket([]byte("history")).Cursor()
for k, v := c.Seek([]byte(from)); k != nil && bytes.Compare(k, []byte(to)) <= 0; k, v = c.Next() {
var l LogLine
if err := json.Unmarshal(v, &l); err != nil {
return err
}
logs = append(logs, l)
}
return nil
})
return logs, err
}
// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
b := make([]byte, 8)
binary.BigEndian.PutUint64(b, v)
return b
}