apiary/honeypot/store/store_test.go
2021-09-17 02:01:43 +02:00

98 lines
2.1 KiB
Go

package store_test
import (
"math/rand"
"net"
"testing"
"time"
"github.com/google/uuid"
"github.uio.no/torjus/apiary/honeypot/store"
"github.uio.no/torjus/apiary/models"
)
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)
}
if len(all) != len(testAttempts) {
t.Errorf("All returned wrong amount. Got %d want %d", len(all), 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))
}
}
})
}
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"
}
}