package store import ( "fmt" "sort" "strings" "sync" "git.t-juice.club/torjus/apiary/models" ) var _ LoginAttemptStore = &MemoryStore{} type MemoryStore struct { lock sync.RWMutex attempts []models.LoginAttempt currentID int } type StatItem struct { Key string Count int } type StatItems []StatItem func (ms *MemoryStore) AddAttempt(l *models.LoginAttempt) error { ms.lock.Lock() defer ms.lock.Unlock() l.ID = ms.currentID + 1 ms.currentID = ms.currentID + 1 ms.attempts = append(ms.attempts, *l) return nil } func (ms *MemoryStore) All() (<-chan models.LoginAttempt, error) { ch := make(chan models.LoginAttempt) go func() { ms.lock.RLock() defer ms.lock.RUnlock() for _, attempt := range ms.attempts { ch <- attempt } close(ch) }() return ch, nil } func (ms *MemoryStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) { counts := make(map[string]int) if statType == LoginStatsTotals { return ms.statTotals() } ms.lock.RLock() defer ms.lock.RUnlock() for _, a := range ms.attempts { switch statType { case LoginStatsPasswords: counts[a.Password]++ case LoginStatsCountry: counts[a.Country]++ case LoginStatsIP: counts[a.RemoteIP.String()]++ case LoginStatsUsername: counts[a.Username]++ default: return nil, fmt.Errorf("invalid stat type") } } 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 (ss StatItems) Len() int { return len(ss) } func (ss StatItems) Less(i, j int) bool { return ss[i].Count < ss[j].Count } func (ss StatItems) Swap(i, j int) { ss[i], ss[j] = ss[j], ss[i] } func (ms *MemoryStore) statTotals() ([]StatsResult, error) { passwords := make(map[string]int) usernames := make(map[string]int) ips := make(map[string]int) countries := make(map[string]int) ms.lock.RLock() defer ms.lock.RUnlock() for _, val := range ms.attempts { passwords[val.Password] += 1 usernames[val.Username] += 1 ips[val.RemoteIP.String()] += 1 countries[val.Country] += 1 } 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: len(ms.attempts)}, } return stats, nil } func (ms *MemoryStore) Query(query AttemptQuery) ([]models.LoginAttempt, error) { var results []models.LoginAttempt ms.lock.Lock() defer ms.lock.Unlock() for _, la := range ms.attempts { switch query.QueryType { case AttemptQueryTypeIP: if la.RemoteIP.String() == query.Query { results = append(results, la) } case AttemptQueryTypePassword: if strings.Contains(la.Password, query.Query) { results = append(results, la) } case AttemptQueryTypeUsername: if strings.Contains(la.Username, query.Query) { results = append(results, la) } } } return results, nil } func (s *MemoryStore) IsHealthy() error { return nil } func toResults(m map[string]int) []StatsResult { var results []StatsResult for key, value := range m { results = append(results, StatsResult{key, value}) } return results }