Reorganize packages

honeypot/ssh/store/memory.go

@@ -0,0 +1,153 @@
+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
+}