apiary/honeypot/ssh/store/metrics.go

package store

import (
	"context"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.uio.no/torjus/apiary/models"
)

type MetricsCollectingStore struct {
	store                LoginAttemptStore
	attemptsCounter      *prometheus.CounterVec
	uniqueUsernamesCount prometheus.Gauge
	uniquePasswordsCount prometheus.Gauge
	uniqueIPsCount       prometheus.Gauge
}

func NewMetricsCollectingStore(ctx context.Context, store LoginAttemptStore) *MetricsCollectingStore {
	mcs := &MetricsCollectingStore{store: store}

	mcs.attemptsCounter = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name:        "apiary_ssh_attempt_counter",
			Help:        "Total count of login attempts toward SSH",
			ConstLabels: prometheus.Labels{"service": "honeypot_ssh"},
		},
		[]string{"CountryCode"},
	)
	mcs.uniqueUsernamesCount = prometheus.NewGauge(
		prometheus.GaugeOpts{
			Name:        "apiary_ssh_unique_usernames_count",
			Help:        "Counter of unique usernames.",
			ConstLabels: prometheus.Labels{"service": "honeypot_ssh"},
		})

	mcs.uniquePasswordsCount = prometheus.NewGauge(
		prometheus.GaugeOpts{
			Name:        "apiary_ssh_unique_passwords_count",
			Help:        "Counter of unique passwords.",
			ConstLabels: prometheus.Labels{"service": "honeypot_ssh"},
		})

	mcs.uniqueIPsCount = prometheus.NewGauge(
		prometheus.GaugeOpts{
			Name:        "apiary_ssh_unique_ips_count",
			Help:        "Counter of unique IPs.",
			ConstLabels: prometheus.Labels{"service": "honeypot_ssh"},
		})

	prometheus.MustRegister(mcs.attemptsCounter)
	prometheus.MustRegister(mcs.uniqueUsernamesCount)
	prometheus.MustRegister(mcs.uniquePasswordsCount)
	prometheus.MustRegister(mcs.uniqueIPsCount)

	// Kinda jank, we just fetch the stats every 10seconds, but it should be cached most of the time.
	go func(ctx context.Context) {
		ticker := time.NewTicker(10 * time.Second)
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			mcs.Stats(LoginStatsTotals, 0)
		}
	}(ctx)

	return mcs
}

func (s *MetricsCollectingStore) AddAttempt(l *models.LoginAttempt) error {
	err := s.store.AddAttempt(l)
	s.attemptsCounter.WithLabelValues(l.Country).Inc()
	return err
}

func (s *MetricsCollectingStore) All() ([]models.LoginAttempt, error) {
	return s.store.All()
}

func (s *MetricsCollectingStore) Stats(statType LoginStats, limit int) ([]StatsResult, error) {
	stats, err := s.store.Stats(statType, limit)
	if statType == LoginStatsTotals {
		for _, element := range stats {
			switch element.Name {
			case "UniquePasswords":
				s.uniquePasswordsCount.Set(float64(element.Count))
			case "UniqueUsernames":
				s.uniqueUsernamesCount.Set(float64(element.Count))
			case "UniqueIPs":
				s.uniqueIPsCount.Set(float64(element.Count))
			default:
				continue
			}
		}
	}
	return stats, err

}

func (s *MetricsCollectingStore) Query(query AttemptQuery) ([]models.LoginAttempt, error) {
	return s.store.Query(query)
}