154 lines
3.1 KiB
Go
154 lines
3.1 KiB
Go
package store
|
|
|
|
import (
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.uio.no/torjus/apiary/models"
|
|
)
|
|
|
|
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() ([]models.LoginAttempt, error) {
|
|
return ms.attempts, 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 toResults(m map[string]int) []StatsResult {
|
|
var results []StatsResult
|
|
|
|
for key, value := range m {
|
|
results = append(results, StatsResult{key, value})
|
|
}
|
|
|
|
return results
|
|
}
|