apiary/honeypot/ssh/store/store_test.go

package store_test

import (
	"math/rand"
	"net"
	"testing"
	"time"

	"git.t-juice.club/torjus/apiary/honeypot/ssh/store"
	"git.t-juice.club/torjus/apiary/models"
	"github.com/google/uuid"
)

func testLoginAttemptStore(s store.LoginAttemptStore, t *testing.T) {
	t.Run("Simple", func(t *testing.T) {
		testAttempts := randomAttempts(10)

		for _, attempt := range testAttempts {
			if err := s.AddAttempt(attempt); err != nil {
				t.Fatalf("Error adding attempt: %s", err)
			}
		}

		all, err := s.All()
		if err != nil {
			t.Fatalf("Error getting all attempts: %s", err)
		}
		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 {
			t.Errorf("Stats returned error: %s", err)
		}
		for _, stat := range stats {
			if stat.Name == "TotalLoginAttempts" && stat.Count != len(testAttempts) {
				t.Errorf("Stats for total attempts is wrong. Got %d want %d", stat.Count, len(testAttempts))
			}
		}
	})
	t.Run("Query", func(t *testing.T) {
		testAttempts := []*models.LoginAttempt{
			{
				Date:     time.Now(),
				RemoteIP: net.ParseIP("127.0.0.1"),
				Username: "corndog",
				Password: "corndog",
			},
			{
				Date:     time.Now(),
				RemoteIP: net.ParseIP("127.0.0.1"),
				Username: "corndog",
				Password: "c0rnd0g",
			},
			{
				Date:     time.Now(),
				RemoteIP: net.ParseIP("10.0.0.1"),
				Username: "root",
				Password: "password",
			},
			{
				Date:     time.Now(),
				RemoteIP: net.ParseIP("10.0.0.2"),
				Username: "ubnt",
				Password: "password",
			},
		}

		for _, attempt := range testAttempts {
			err := s.AddAttempt(attempt)
			if err != nil {
				t.Fatalf("Unable to add attempt: %s", err)
			}
		}
		testCases := []struct {
			Name           string
			Query          store.AttemptQuery
			ExpectedResult []models.LoginAttempt
		}{
			{
				Name:  "password one result",
				Query: store.AttemptQuery{QueryType: store.AttemptQueryTypePassword, Query: "corndog"},
				ExpectedResult: []models.LoginAttempt{
					{
						RemoteIP: net.ParseIP("127.0.0.1"),
						Username: "corndog",
						Password: "corndog",
					},
				},
			},
			{
				Name:  "username one result",
				Query: store.AttemptQuery{QueryType: store.AttemptQueryTypeUsername, Query: "root"},
				ExpectedResult: []models.LoginAttempt{
					{
						RemoteIP: net.ParseIP("10.0.0.1"),
						Username: "root",
						Password: "password",
					},
				},
			},
			{
				Name:  "username two results",
				Query: store.AttemptQuery{QueryType: store.AttemptQueryTypeUsername, Query: "corndog"},
				ExpectedResult: []models.LoginAttempt{
					{
						RemoteIP: net.ParseIP("127.0.0.1"),
						Username: "corndog",
						Password: "c0rnd0g",
					},
					{
						RemoteIP: net.ParseIP("127.0.0.1"),
						Username: "corndog",
						Password: "corndog",
					},
				},
			},
		}

		for _, tc := range testCases {
			res, err := s.Query(tc.Query)
			if err != nil {
				t.Errorf("Error performing query: %s", err)
			}
			if !equalAttempts(res, tc.ExpectedResult) {
				t.Errorf("Query did not return expected results")
				t.Logf("%+v", res)
				t.Logf("%+v", tc.ExpectedResult)

			}
		}
	})

	t.Run("QueryCache", func(t *testing.T) {
		err := s.AddAttempt(&models.LoginAttempt{RemoteIP: net.ParseIP("127.0.0.1"), Username: "test", Password: "test"})
		if err != nil {
			t.Fatalf("Error adding attempt: %s", err)
		}

		res, err := s.Query(store.AttemptQuery{QueryType: store.AttemptQueryTypeUsername, Query: "test"})
		if err != nil {
			t.Fatalf("Error adding attempt: %s", err)
		}
		if len(res) != 1 {
			t.Errorf("Wrong amount of results")
		}

		err = s.AddAttempt(&models.LoginAttempt{RemoteIP: net.ParseIP("127.0.0.1"), Username: "test", Password: "best"})
		if err != nil {
			t.Fatalf("Error adding attempt: %s", err)
		}
		res, err = s.Query(store.AttemptQuery{QueryType: store.AttemptQueryTypeUsername, Query: "test"})
		if err != nil {
			t.Fatalf("Error adding attempt: %s", err)
		}
		if len(res) != 2 {
			t.Errorf("Wrong amount of results")
		}
	})
	t.Run("QueryStats", func(t *testing.T) {
		firstStats, err := s.Stats(store.LoginStatsTotals, 1)
		if err != nil {
			t.Fatalf("Error getting stats: %s", err)
		}
		err = s.AddAttempt(&models.LoginAttempt{RemoteIP: net.ParseIP("127.0.0.1"), Username: "test", Password: "best"})
		if err != nil {
			t.Fatalf("Error adding attempt: %s", err)
		}
		secondStats, err := s.Stats(store.LoginStatsTotals, 1)
		if err != nil {
			t.Fatalf("Error getting stats: %s", err)
		}
		var firstCount, secondCount int
		for _, stat := range firstStats {
			if stat.Name == "TotalLoginAttempts" {
				firstCount = stat.Count
			}
		}
		for _, stat := range secondStats {
			if stat.Name == "TotalLoginAttempts" {
				secondCount = stat.Count
			}
		}

		if secondCount != firstCount+1 {
			t.Errorf("TotalLoginAttempts did not increment")
		}
	})
}

func benchmarkLoginAttemptStore(setupFunc func() store.LoginAttemptStore, b *testing.B) {
	b.Run("BenchmarkAdd", func(b *testing.B) {
		s := setupFunc()
		for i := 0; i < b.N; i++ {
			attempt := randomAttempts(1)
			err := s.AddAttempt(attempt[0])
			if err != nil {
				b.Fatalf("Error adding attempt: %s", err)
			}
		}
	})
	b.Run("BenchmarkAdd10k", func(b *testing.B) {
		attempts := randomAttempts(10_000)
		for i := 0; i < b.N; i++ {
			b.StopTimer()
			s := setupFunc()
			b.StartTimer()
			for _, attempt := range attempts {
				err := s.AddAttempt(attempt)
				if err != nil {
					b.Fatalf("Error adding attempt: %s", err)
				}
			}
		}
	})
	b.Run("BenchmarkAll10k", func(b *testing.B) {
		s := setupFunc()
		attempts := randomAttempts(10_000)
		for _, attempt := range attempts {
			err := s.AddAttempt(attempt)
			if err != nil {
				b.Fatalf("Error adding attempt: %s", err)
			}
		}

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			all, err := s.All()
			if err != nil {
				b.Fatalf("Error fetchin all: %s", err)
			}
			var count int
			for range all {
				count++
			}
			_ = count
		}
	})
}

func randomAttempts(count int) []*models.LoginAttempt {
	var attempts []*models.LoginAttempt
	for i := 0; i < count; i++ {
		attempt := &models.LoginAttempt{
			Date:             time.Now(),
			RemoteIP:         randomIP(),
			Username:         randomString(8),
			Password:         randomString(8),
			Country:          randomCountry(),
			ConnectionUUID:   uuid.Must(uuid.NewRandom()),
			SSHClientVersion: "SSH TEST LOL",
		}
		attempts = append(attempts, attempt)
	}
	return attempts
}

func randomIP() net.IP {
	a := byte(rand.Intn(254))
	b := byte(rand.Intn(254))
	c := byte(rand.Intn(254))
	d := byte(rand.Intn(254))
	return net.IPv4(a, b, c, d)
}

func randomString(n int) string {
	const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

func randomCountry() string {
	switch rand.Intn(10) {
	case 1:
		return "CN"
	case 2:
		return "US"
	case 3:
		return "NO"
	case 4:
		return "RU"
	case 5:
		return "DE"
	case 6:
		return "FI"
	case 7:
		return "BR"
	default:
		return "SE"
	}
}

func equalAttempts(a, b []models.LoginAttempt) bool {
	if len(a) != len(b) {
		return false
	}

	aFound := make([]bool, len(a))

	for i, aAttempt := range a {
		for _, bAttempt := range b {
			if aAttempt.Username == bAttempt.Username &&
				aAttempt.Password == bAttempt.Password &&
				aAttempt.RemoteIP.String() == bAttempt.RemoteIP.String() {
				aFound[i] = true
			}
		}
	}

	for _, found := range aFound {
		if !found {
			return false
		}
	}

	return true
}

func fuzzLoginAttemptStore(s store.LoginAttemptStore, f *testing.F) {
	usernames := []string{"username", "root", "ubnt", "pi", "admin"}
	for _, username := range usernames {
		f.Add(username)
	}
	f.Fuzz(func(t *testing.T, orig string) {
		attempt := models.LoginAttempt{
			Date:     time.Now(),
			Username: orig,
			Password: randomString(8),
			Country:  "NO",
		}

		if err := s.AddAttempt(&attempt); err != nil {
			t.Fatalf("error adding: %s", err)
		}
	})
}