package store import ( "encoding/binary" "encoding/json" "fmt" "sort" "strings" "git.t-juice.club/torjus/apiary/models" bolt "go.etcd.io/bbolt" ) // var _ LoginAttemptStore = &BBoltStore{} var bktKeyLogins []byte = []byte("logins") type BBoltStore struct { db *bolt.DB } func NewBBoltStore(path string) (*BBoltStore, error) { db, err := bolt.Open(path, 0o666, nil) if err != nil { return nil, fmt.Errorf("error opening database: %w", err) } var store BBoltStore store.db = db err = db.Update(func(tx *bolt.Tx) error { _, err := tx.CreateBucketIfNotExists(bktKeyLogins) return err }) if err != nil { return nil, fmt.Errorf("error creating database bucket: %w", err) } return &store, nil } func (s *BBoltStore) Close() error { return s.db.Close() } func (s *BBoltStore) AddAttempt(l *models.LoginAttempt) error { data, err := json.Marshal(l) if err != nil { return err } return s.db.Update(func(tx *bolt.Tx) error { bkt := tx.Bucket(bktKeyLogins) seq, err := bkt.NextSequence() if err != nil { return err } key := itob(seq) return bkt.Put(key, data) }) } func (s *BBoltStore) All() (<-chan models.LoginAttempt, error) { ch := make(chan models.LoginAttempt) go func() { _ = s.db.View(func(tx *bolt.Tx) error { bkt := tx.Bucket(bktKeyLogins) c := bkt.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { var l models.LoginAttempt if err := json.Unmarshal(v, &l); err != nil { close(ch) panic(err) } ch <- l } close(ch) return nil }) }() return ch, nil } func (s *BBoltStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) { if statType == LoginStatsTotals { return s.statTotals() } counts := make(map[string]int) err := s.db.View(func(tx *bolt.Tx) error { bkt := tx.Bucket(bktKeyLogins) c := bkt.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { var l models.LoginAttempt if err := json.Unmarshal(v, &l); err != nil { return err } switch statType { case LoginStatsPasswords: counts[l.Password]++ case LoginStatsCountry: counts[l.Country]++ case LoginStatsIP: counts[l.RemoteIP.String()]++ case LoginStatsUsername: counts[l.Username]++ default: return fmt.Errorf("invalid stat type") } } return nil }) if err != nil { return nil, fmt.Errorf("error generating stats: %w", err) } if limit < 1 { return toResults(counts), nil } if limit >= len(counts) { return toResults(counts), nil } var si StatItems for key := range counts { si = append(si, StatItem{Key: key, Count: counts[key]}) } sort.Sort(si) output := make(map[string]int) for i := len(si) - 1; i > len(si)-limit-1; i-- { output[si[i].Key] = si[i].Count } return toResults(output), nil } func (s *BBoltStore) statTotals() ([]StatsResult, error) { passwords := make(map[string]int) usernames := make(map[string]int) ips := make(map[string]int) countries := make(map[string]int) var count int err := s.db.View(func(tx *bolt.Tx) error { bkt := tx.Bucket(bktKeyLogins) c := bkt.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { var l models.LoginAttempt if err := json.Unmarshal(v, &l); err != nil { return err } passwords[l.Password] += 1 usernames[l.Username] += 1 ips[l.RemoteIP.String()] += 1 countries[l.Country] += 1 count++ } return nil }) if err != nil { return nil, err } stats := []StatsResult{ {Name: "UniquePasswords", Count: len(passwords)}, {Name: "UniqueUsernames", Count: len(usernames)}, {Name: "UniqueIPs", Count: len(ips)}, {Name: "UniqueCountries", Count: len(countries)}, {Name: "TotalLoginAttempts", Count: count}, } return stats, nil } func (s *BBoltStore) Query(query AttemptQuery) ([]models.LoginAttempt, error) { var results []models.LoginAttempt err := s.db.View(func(tx *bolt.Tx) error { bkt := tx.Bucket(bktKeyLogins) c := bkt.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { var l models.LoginAttempt if err := json.Unmarshal(v, &l); err != nil { return err } switch query.QueryType { case AttemptQueryTypeIP: if l.RemoteIP.String() == query.Query { results = append(results, l) } case AttemptQueryTypePassword: if strings.Contains(l.Password, query.Query) { results = append(results, l) } case AttemptQueryTypeUsername: if strings.Contains(l.Username, query.Query) { results = append(results, l) } } } return nil }) if err != nil { return nil, err } return results, nil } func (s *BBoltStore) IsHealthy() error { // TODO: Do actual healthcheck return nil } func itob(v uint64) []byte { b := make([]byte, 8) binary.BigEndian.PutUint64(b, v) return b }