Cache total stats in memory

honeypot/ssh/store/cache.go

@@ -9,15 +9,40 @@ type CachingStore struct {
 	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 {
-	return &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)
+	}
+
+	cs.totalLoginsCount = len(all)
+	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{}{}
+	}
+
+	return cs
 }
 
 func (s *CachingStore) AddAttempt(l *models.LoginAttempt) error {
@@ -25,6 +50,11 @@ func (s *CachingStore) AddAttempt(l *models.LoginAttempt) error {
 	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)
 }
 
@@ -35,20 +65,15 @@ func (s *CachingStore) All() ([]models.LoginAttempt, error) {
 func (s *CachingStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {
 	// Only cache totals for now, as they are the most queried
 	if statType == LoginStatsTotals {
-		if s.totalStatsCache != nil {
-			return s.totalStatsCache, nil
-		}
-		stats, err := s.backend.Stats(statType, limit)
-		if err != nil {
-			return stats, err
-		}
-		s.totalStatsCache = stats
-
-		return stats, err
+		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
 	}
 
-	// TODO: Maybe cache the default limits
-
 	return s.backend.Stats(statType, limit)
 }