apiary/honeypot/ssh/store/metrics.go

package store

import (
	"context"
	"time"

	"git.t-juice.club/torjus/apiary/models"
	"github.com/prometheus/client_golang/prometheus"
)

var _ LoginAttemptStore = &MetricsCollectingStore{}

const tickDuration = 5 * time.Second

type MetricsCollectingStore struct {
	store LoginAttemptStore

	attemptsCounter      *prometheus.CounterVec
	uniqueUsernamesCount prometheus.Gauge
	uniquePasswordsCount prometheus.Gauge
	uniqueIPsCount       prometheus.Gauge
	totalAttemptsCount   prometheus.Gauge

	statsTicker *time.Ticker
}

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"},
		})

	mcs.totalAttemptsCount = prometheus.NewGauge(
		prometheus.GaugeOpts{
			Name:        "apiary_ssh_total_login_attempts",
			Help:        "Total amount of login attempts made.",
			ConstLabels: prometheus.Labels{"service": "honeypot_ssh"},
		})

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

	mcs.statsTicker = time.NewTicker(tickDuration)
	go func() {
		defer mcs.statsTicker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-mcs.statsTicker.C:
				mcs.Stats(LoginStatsTotals, 0) // nolint: errcheck
			}
		}
	}()

	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() (<-chan 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 {
		// Reset ticker to avoid updating twice within the duration
		s.statsTicker.Reset(tickDuration)
		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))
			case "TotalLoginAttempts":
				s.totalAttemptsCount.Set(float64(element.Count))
			default:
				continue
			}
		}
	}
	return stats, err
}

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

func (s *MetricsCollectingStore) IsHealthy() error {
	return s.store.IsHealthy()
}