apiary/honeypot/ssh/store/memory.go

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
}