apiary/honeypot/ssh/store/cache.go

package store

import "git.t-juice.club/torjus/apiary/models"

var _ LoginAttemptStore = &CachingStore{}

type CachingStore struct {
	backend LoginAttemptStore

	totalStatsCache    []StatsResult
	usernameQueryCache map[string][]models.LoginAttempt
	passwordQueryCache map[string][]models.LoginAttempt
	ipQueryCache       map[string][]models.LoginAttempt
	uniqueUsernames    map[string]struct{}
	uniquePasswords    map[string]struct{}
	uniqueIPs          map[string]struct{}
	uniqueCountries    map[string]struct{}
	totalLoginsCount   int
}

func NewCachingStore(backend LoginAttemptStore) *CachingStore {
	cs := &CachingStore{
		backend:            backend,
		usernameQueryCache: make(map[string][]models.LoginAttempt),
		passwordQueryCache: make(map[string][]models.LoginAttempt),
		ipQueryCache:       make(map[string][]models.LoginAttempt),
		uniqueUsernames:    make(map[string]struct{}),
		uniquePasswords:    make(map[string]struct{}),
		uniqueIPs:          make(map[string]struct{}),
		uniqueCountries:    make(map[string]struct{}),
	}

	all, err := backend.All()
	if err != nil {
		//TODO: Handle better maybe?
		panic(err)
	}
	var loginCount int

	for attempt := range all {
		cs.uniqueUsernames[attempt.Username] = struct{}{}
		cs.uniquePasswords[attempt.Password] = struct{}{}
		cs.uniqueIPs[attempt.RemoteIP.String()] = struct{}{}
		cs.uniqueCountries[attempt.Country] = struct{}{}
		loginCount++
	}
	cs.totalLoginsCount = loginCount

	return cs
}

func (s *CachingStore) AddAttempt(l *models.LoginAttempt) error {
	s.totalStatsCache = nil
	delete(s.ipQueryCache, l.RemoteIP.String())
	delete(s.passwordQueryCache, l.Password)
	delete(s.usernameQueryCache, l.Username)
	s.totalLoginsCount++
	s.uniqueUsernames[l.Username] = struct{}{}
	s.uniquePasswords[l.Password] = struct{}{}
	s.uniqueIPs[l.RemoteIP.String()] = struct{}{}
	s.uniqueCountries[l.Country] = struct{}{}
	return s.backend.AddAttempt(l)
}

func (s *CachingStore) All() (<-chan models.LoginAttempt, error) {
	return s.backend.All()
}

func (s *CachingStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {
	// Only cache totals for now, as they are the most queried
	if statType == LoginStatsTotals {
		return []StatsResult{
			{Name: "UniquePasswords", Count: len(s.uniquePasswords)},
			{Name: "UniqueUsernames", Count: len(s.uniqueUsernames)},
			{Name: "UniqueIPs", Count: len(s.uniqueIPs)},
			{Name: "UniqueCountries", Count: len(s.uniqueCountries)},
			{Name: "TotalLoginAttempts", Count: s.totalLoginsCount},
		}, nil
	}

	return s.backend.Stats(statType, limit)
}

func (s *CachingStore) Query(query AttemptQuery) ([]models.LoginAttempt, error) {
	switch query.QueryType {
	case AttemptQueryTypeIP:
		if attempts, ok := s.ipQueryCache[query.Query]; ok {
			return attempts, nil
		}
	case AttemptQueryTypePassword:
		if attempts, ok := s.passwordQueryCache[query.Query]; ok {
			return attempts, nil
		}
	case AttemptQueryTypeUsername:
		if attempts, ok := s.usernameQueryCache[query.Query]; ok {
			return attempts, nil
		}
	}
	attempts, err := s.backend.Query(query)
	if err != nil {
		return attempts, err
	}

	switch query.QueryType {
	case AttemptQueryTypeIP:
		s.ipQueryCache[query.Query] = attempts
	case AttemptQueryTypeUsername:
		s.ipQueryCache[query.Query] = attempts
	case AttemptQueryTypePassword:
		s.ipQueryCache[query.Query] = attempts
	}

	return attempts, err
}

func (s *CachingStore) IsHealthy() error {
	return s.backend.IsHealthy()
}