diff --git a/honeypot/ssh/store/cache.go b/honeypot/ssh/store/cache.go index 39b6f5d..a5ef880 100644 --- a/honeypot/ssh/store/cache.go +++ b/honeypot/ssh/store/cache.go @@ -2,6 +2,8 @@ package store import "git.t-juice.club/torjus/apiary/models" +var _ LoginAttemptStore = &CachingStore{} + type CachingStore struct { backend LoginAttemptStore @@ -33,14 +35,16 @@ func NewCachingStore(backend LoginAttemptStore) *CachingStore { //TODO: Handle better maybe? panic(err) } + var loginCount int - cs.totalLoginsCount = len(all) - for _, attempt := range 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{}{} + loginCount++ } + cs.totalLoginsCount = loginCount return cs } @@ -58,7 +62,7 @@ func (s *CachingStore) AddAttempt(l *models.LoginAttempt) error { return s.backend.AddAttempt(l) } -func (s *CachingStore) All() ([]models.LoginAttempt, error) { +func (s *CachingStore) All() (<-chan models.LoginAttempt, error) { return s.backend.All() } diff --git a/honeypot/ssh/store/memory.go b/honeypot/ssh/store/memory.go index 4bae0ed..e064ae5 100644 --- a/honeypot/ssh/store/memory.go +++ b/honeypot/ssh/store/memory.go @@ -9,6 +9,8 @@ import ( "git.t-juice.club/torjus/apiary/models" ) +var _ LoginAttemptStore = &MemoryStore{} + type MemoryStore struct { lock sync.RWMutex attempts []models.LoginAttempt @@ -32,8 +34,17 @@ func (ms *MemoryStore) AddAttempt(l *models.LoginAttempt) error { return nil } -func (ms *MemoryStore) All() ([]models.LoginAttempt, error) { - return ms.attempts, 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) { diff --git a/honeypot/ssh/store/metrics.go b/honeypot/ssh/store/metrics.go index 34296f3..509408f 100644 --- a/honeypot/ssh/store/metrics.go +++ b/honeypot/ssh/store/metrics.go @@ -8,6 +8,8 @@ import ( "github.com/prometheus/client_golang/prometheus" ) +var _ LoginAttemptStore = &MetricsCollectingStore{} + const tickDuration = 5 * time.Second type MetricsCollectingStore struct { @@ -89,7 +91,7 @@ func (s *MetricsCollectingStore) AddAttempt(l *models.LoginAttempt) error { return err } -func (s *MetricsCollectingStore) All() ([]models.LoginAttempt, error) { +func (s *MetricsCollectingStore) All() (<-chan models.LoginAttempt, error) { return s.store.All() } diff --git a/honeypot/ssh/store/postgres.go b/honeypot/ssh/store/postgres.go index 5a8c311..0376f76 100644 --- a/honeypot/ssh/store/postgres.go +++ b/honeypot/ssh/store/postgres.go @@ -9,6 +9,8 @@ import ( _ "github.com/jackc/pgx/v4/stdlib" ) +var _ LoginAttemptStore = &PostgresStore{} + const DBSchema = ` CREATE TABLE IF NOT EXISTS login_attempts( id serial PRIMARY KEY, @@ -61,27 +63,30 @@ RETURNING id;` return tx.Commit() } -func (s *PostgresStore) All() ([]models.LoginAttempt, error) { +func (s *PostgresStore) All() (<-chan models.LoginAttempt, error) { stmt := `SELECT date, remote_ip, username, password, client_version, connection_uuid, country FROM login_attempts` rows, err := s.db.Query(stmt) if err != nil { return nil, err } - defer rows.Close() - var attempts []models.LoginAttempt - for rows.Next() { - var a models.LoginAttempt - var ip string - if err := rows.Scan(&a.Date, &ip, &a.Username, &a.Password, &a.SSHClientVersion, &a.SSHClientVersion, &a.Country); err != nil { - return nil, err + ch := make(chan models.LoginAttempt) + go func() { + defer rows.Close() + for rows.Next() { + var a models.LoginAttempt + var ip string + if err := rows.Scan(&a.Date, &ip, &a.Username, &a.Password, &a.SSHClientVersion, &a.SSHClientVersion, &a.Country); err != nil { + panic(err) + } + a.RemoteIP = net.ParseIP(ip) + ch <- a } - a.RemoteIP = net.ParseIP(ip) - attempts = append(attempts, a) - } + close(ch) + }() - return attempts, nil + return ch, nil } func (s *PostgresStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) { diff --git a/honeypot/ssh/store/store.go b/honeypot/ssh/store/store.go index 1c091bc..f71b641 100644 --- a/honeypot/ssh/store/store.go +++ b/honeypot/ssh/store/store.go @@ -38,7 +38,7 @@ type AttemptQuery struct { } type LoginAttemptStore interface { AddAttempt(l *models.LoginAttempt) error - All() ([]models.LoginAttempt, error) + All() (<-chan models.LoginAttempt, error) Stats(statType LoginStats, limit int) ([]StatsResult, error) Query(query AttemptQuery) ([]models.LoginAttempt, error) HealthCheker diff --git a/honeypot/ssh/store/store_test.go b/honeypot/ssh/store/store_test.go index abf6c9e..2f140eb 100644 --- a/honeypot/ssh/store/store_test.go +++ b/honeypot/ssh/store/store_test.go @@ -25,8 +25,12 @@ func testLoginAttemptStore(s store.LoginAttemptStore, t *testing.T) { if err != nil { t.Fatalf("Error getting all attempts: %s", err) } - if len(all) != len(testAttempts) { - t.Errorf("All returned wrong amount. Got %d want %d", len(all), len(testAttempts)) + var count int + for range all { + count++ + } + if count != len(testAttempts) { + t.Errorf("All returned wrong amount. Got %d want %d", count, len(testAttempts)) } stats, err := s.Stats(store.LoginStatsTotals, 1) if err != nil { @@ -229,7 +233,11 @@ func benchmarkLoginAttemptStore(setupFunc func() store.LoginAttemptStore, b *tes if err != nil { b.Fatalf("Error fetchin all: %s", err) } - _ = len(all) + var count int + for range all { + count++ + } + _ = count } }) }