apiary/honeypot/ssh/store/bbolt.go

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
}