119 lines
3.4 KiB
Go
119 lines
3.4 KiB
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()
|
|
}
|