387 lines
12 KiB
Go
387 lines
12 KiB
Go
package database
|
|
|
|
import (
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"fmt"
|
|
"net"
|
|
"reflect"
|
|
"regexp"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/DATA-DOG/go-sqlmock"
|
|
"github.com/crowdsecurity/crowdsec/pkg/types"
|
|
"github.com/jinzhu/gorm"
|
|
"github.com/onsi/ginkgo"
|
|
"github.com/onsi/gomega"
|
|
)
|
|
|
|
type AnyTime struct{}
|
|
|
|
// Match satisfies sqlmock.Argument interface
|
|
func (a AnyTime) Match(v driver.Value) bool {
|
|
_, ok := v.(time.Time)
|
|
return ok
|
|
}
|
|
|
|
var _ = ginkgo.Describe("TestWrites", func() {
|
|
var ctx *Context
|
|
var mock sqlmock.Sqlmock
|
|
|
|
ginkgo.BeforeEach(func() {
|
|
var db *sql.DB
|
|
var err error
|
|
|
|
db, mock, err = sqlmock.New() // mock sql.DB
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
|
|
gdb, err := gorm.Open("sqlite", db) // open gorm db
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
|
|
ctx = &Context{Db: gdb}
|
|
//ctx.Db.LogMode(true)
|
|
})
|
|
ginkgo.AfterEach(func() {
|
|
err := mock.ExpectationsWereMet() // make sure all expectations were met
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
})
|
|
|
|
ginkgo.Context("insert ban_applications", func() {
|
|
ginkgo.It("insert 1.2.3.4", func() {
|
|
|
|
const sqlSelectAll = `SELECT * FROM "ban_applications" WHERE "ban_applications"."deleted_at" IS NULL AND (("ban_applications"."ip_text" = ?)) ORDER BY "ban_applications"."id" ASC LIMIT 1`
|
|
|
|
insertBan := types.BanApplication{IpText: "1.2.3.4"}
|
|
|
|
mock.ExpectQuery(regexp.QuoteMeta(sqlSelectAll)).WithArgs("1.2.3.4").WillReturnRows(sqlmock.NewRows(nil))
|
|
|
|
mock.ExpectBegin()
|
|
|
|
const sqlInsertBanApplication = `INSERT INTO "ban_applications" ("created_at","updated_at","deleted_at","measure_source","measure_type","measure_extra","until","start_ip","end_ip","target_cn","target_as","target_as_name","ip_text","reason","scenario","signal_occurence_id") VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`
|
|
InsertExpectedResult := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlInsertBanApplication)).WithArgs(
|
|
AnyTime{},
|
|
AnyTime{},
|
|
nil,
|
|
insertBan.MeasureSource,
|
|
insertBan.MeasureType,
|
|
insertBan.MeasureExtra,
|
|
AnyTime{},
|
|
insertBan.StartIp,
|
|
insertBan.EndIp,
|
|
insertBan.TargetCN,
|
|
insertBan.TargetAS,
|
|
insertBan.TargetASName,
|
|
insertBan.IpText,
|
|
insertBan.Reason,
|
|
insertBan.Scenario,
|
|
insertBan.SignalOccurenceID).WillReturnResult(InsertExpectedResult)
|
|
|
|
mock.ExpectCommit()
|
|
|
|
err := ctx.WriteBanApplication(insertBan)
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
})
|
|
})
|
|
|
|
ginkgo.Context("insert signal_occurence", func() {
|
|
ginkgo.It("insert signal+ban for 1.2.3.4", func() {
|
|
insertBan := types.BanApplication{IpText: "1.2.3.4", SignalOccurenceID: 1}
|
|
insertSig := types.SignalOccurence{
|
|
MapKey: "ratata",
|
|
Scenario: "test_1",
|
|
BanApplications: []types.BanApplication{insertBan},
|
|
Source_ip: "1.2.3.4",
|
|
Source_range: "1.2.3.0/24",
|
|
Source_AutonomousSystemNumber: "1234",
|
|
}
|
|
|
|
//the part that try to delete pending existing bans
|
|
mock.ExpectBegin()
|
|
const sqlDeleteOldBan = `UPDATE "ban_applications" SET "deleted_at"=? WHERE "ban_applications"."deleted_at" IS NULL AND ((ip_text = ?))`
|
|
sqlDeleteOldBanResult := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlDeleteOldBan)).WithArgs(AnyTime{}, "1.2.3.4").WillReturnResult(sqlDeleteOldBanResult)
|
|
mock.ExpectCommit()
|
|
|
|
//insert the signal occurence
|
|
mock.ExpectBegin()
|
|
const sqlInsertNewEvent = `INSERT INTO "signal_occurences" ("created_at","updated_at","deleted_at","map_key","scenario","bucket_id","alert_message","events_count","start_at","stop_at","source_ip","source_range","source_autonomous_system_number","source_autonomous_system_organization","source_country","source_latitude","source_longitude","dest_ip","capacity","leak_speed","reprocess") VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`
|
|
sqlInsertNewEventResult := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlInsertNewEvent)).WithArgs(
|
|
AnyTime{},
|
|
AnyTime{},
|
|
nil,
|
|
insertSig.MapKey,
|
|
insertSig.Scenario,
|
|
"",
|
|
"",
|
|
0,
|
|
AnyTime{},
|
|
AnyTime{},
|
|
insertSig.Source_ip,
|
|
insertSig.Source_range,
|
|
insertSig.Source_AutonomousSystemNumber,
|
|
"",
|
|
"",
|
|
0.0,
|
|
0.0,
|
|
"",
|
|
0,
|
|
0,
|
|
false,
|
|
).WillReturnResult(sqlInsertNewEventResult)
|
|
|
|
//insert the ban application
|
|
const sqlInsertBanApplication = `INSERT INTO "ban_applications" ("created_at","updated_at","deleted_at","measure_source","measure_type","measure_extra","until","start_ip","end_ip","target_cn","target_as","target_as_name","ip_text","reason","scenario","signal_occurence_id") VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`
|
|
sqlInsertBanApplicationResults := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlInsertBanApplication)).WithArgs(
|
|
AnyTime{},
|
|
AnyTime{},
|
|
nil,
|
|
insertBan.MeasureSource,
|
|
insertBan.MeasureType,
|
|
insertBan.MeasureExtra,
|
|
AnyTime{},
|
|
insertBan.StartIp,
|
|
insertBan.EndIp,
|
|
insertBan.TargetCN,
|
|
insertBan.TargetAS,
|
|
insertBan.TargetASName,
|
|
insertBan.IpText,
|
|
insertBan.Reason,
|
|
insertBan.Scenario,
|
|
insertBan.SignalOccurenceID).WillReturnResult(sqlInsertBanApplicationResults)
|
|
|
|
mock.ExpectCommit()
|
|
|
|
err := ctx.WriteSignal(insertSig)
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
})
|
|
})
|
|
|
|
ginkgo.Context("insert old signal_occurence + cleanup", func() {
|
|
ginkgo.It("insert signal+ban for 1.2.3.4", func() {
|
|
|
|
target_ip := net.ParseIP("1.2.3.4")
|
|
|
|
OldBan := types.BanApplication{
|
|
MeasureType: "ban",
|
|
MeasureSource: "local",
|
|
//expired one month ago
|
|
Until: time.Now().Add(-24 * 30 * time.Hour),
|
|
StartIp: types.IP2Int(target_ip),
|
|
EndIp: types.IP2Int(target_ip),
|
|
TargetCN: "FR",
|
|
TargetAS: 1234,
|
|
TargetASName: "Random AS",
|
|
IpText: target_ip.String(),
|
|
Reason: "A reason",
|
|
Scenario: "A scenario",
|
|
}
|
|
OldSignal := types.SignalOccurence{
|
|
MapKey: "lala",
|
|
Scenario: "old_overflow",
|
|
//two month ago : 24*60
|
|
Start_at: time.Now().Add(-24 * 60 * time.Hour),
|
|
Stop_at: time.Now().Add(-24 * 60 * time.Hour),
|
|
BanApplications: []types.BanApplication{OldBan},
|
|
}
|
|
|
|
//the part that try to delete pending existing bans
|
|
mock.ExpectBegin()
|
|
const sqlDeleteOldBan = `UPDATE "ban_applications" SET "deleted_at"=? WHERE "ban_applications"."deleted_at" IS NULL AND ((ip_text = ?))`
|
|
sqlDeleteOldBanResult := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlDeleteOldBan)).WithArgs(AnyTime{}, target_ip.String()).WillReturnResult(sqlDeleteOldBanResult)
|
|
mock.ExpectCommit()
|
|
|
|
//insert the signal occurence
|
|
mock.ExpectBegin()
|
|
const sqlInsertNewEvent = `INSERT INTO "signal_occurences" ("created_at","updated_at","deleted_at","map_key","scenario","bucket_id","alert_message","events_count","start_at","stop_at","source_ip","source_range","source_autonomous_system_number","source_autonomous_system_organization","source_country","source_latitude","source_longitude","dest_ip","capacity","leak_speed","reprocess") VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`
|
|
sqlInsertNewEventResult := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlInsertNewEvent)).WithArgs(
|
|
AnyTime{},
|
|
AnyTime{},
|
|
nil,
|
|
OldSignal.MapKey,
|
|
OldSignal.Scenario,
|
|
"",
|
|
"",
|
|
0,
|
|
AnyTime{},
|
|
AnyTime{},
|
|
OldSignal.Source_ip,
|
|
OldSignal.Source_range,
|
|
OldSignal.Source_AutonomousSystemNumber,
|
|
"",
|
|
"",
|
|
0.0,
|
|
0.0,
|
|
"",
|
|
0,
|
|
0,
|
|
false,
|
|
).WillReturnResult(sqlInsertNewEventResult)
|
|
|
|
//insert the ban application
|
|
const sqlInsertBanApplication = `INSERT INTO "ban_applications" ("created_at","updated_at","deleted_at","measure_source","measure_type","measure_extra","until","start_ip","end_ip","target_cn","target_as","target_as_name","ip_text","reason","scenario","signal_occurence_id") VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`
|
|
sqlInsertBanApplicationResults := sqlmock.NewResult(1, 1)
|
|
mock.ExpectExec(regexp.QuoteMeta(sqlInsertBanApplication)).WithArgs(
|
|
AnyTime{},
|
|
AnyTime{},
|
|
nil,
|
|
OldBan.MeasureSource,
|
|
OldBan.MeasureType,
|
|
OldBan.MeasureExtra,
|
|
AnyTime{},
|
|
OldBan.StartIp,
|
|
OldBan.EndIp,
|
|
OldBan.TargetCN,
|
|
OldBan.TargetAS,
|
|
OldBan.TargetASName,
|
|
OldBan.IpText,
|
|
OldBan.Reason,
|
|
OldBan.Scenario,
|
|
1).WillReturnResult(sqlInsertBanApplicationResults)
|
|
|
|
mock.ExpectCommit()
|
|
|
|
err := ctx.WriteSignal(OldSignal)
|
|
gomega.Expect(err).ShouldNot(gomega.HaveOccurred())
|
|
|
|
})
|
|
})
|
|
|
|
})
|
|
|
|
func TestInsertSqlMock(t *testing.T) {
|
|
gomega.RegisterFailHandler(ginkgo.Fail)
|
|
ginkgo.RunSpecs(t, "TestWrites")
|
|
}
|
|
|
|
func TestInsertOldBans(t *testing.T) {
|
|
//Plan:
|
|
// - flush db
|
|
// - insert month old ban
|
|
// - use GetBansAt on current + past time and check results
|
|
// - @todo : we need to call the DeleteExpired and such
|
|
|
|
validCfg := map[string]string{
|
|
"type": "sqlite",
|
|
"db_path": "./test.db",
|
|
"max_records": "1000",
|
|
"max_records_age": "72h",
|
|
"debug": "false",
|
|
"flush": "true",
|
|
}
|
|
ctx, err := NewDatabase(validCfg)
|
|
if err != nil || ctx == nil {
|
|
t.Fatalf("failed to create simple sqlite")
|
|
}
|
|
|
|
if err := ctx.DeleteAll(); err != nil {
|
|
t.Fatalf("failed to flush existing bans")
|
|
}
|
|
|
|
target_ip := net.ParseIP("1.2.3.4")
|
|
|
|
OldBan := types.BanApplication{
|
|
MeasureType: "ban",
|
|
MeasureSource: "local",
|
|
//expired one month ago
|
|
Until: time.Now().Add(-24 * 30 * time.Hour),
|
|
StartIp: types.IP2Int(target_ip),
|
|
EndIp: types.IP2Int(target_ip),
|
|
TargetCN: "FR",
|
|
TargetAS: 1234,
|
|
TargetASName: "Random AS",
|
|
IpText: target_ip.String(),
|
|
Reason: "A reason",
|
|
Scenario: "A scenario",
|
|
}
|
|
OldSignal := types.SignalOccurence{
|
|
MapKey: "lala",
|
|
Scenario: "old_overflow",
|
|
//two month ago : 24*60
|
|
Start_at: time.Now().Add(-24 * 60 * time.Hour),
|
|
Stop_at: time.Now().Add(-24 * 60 * time.Hour),
|
|
BanApplications: []types.BanApplication{OldBan},
|
|
}
|
|
//write the old signal
|
|
err = ctx.WriteSignal(OldSignal)
|
|
if err != nil {
|
|
t.Fatalf("Failed to insert old signal : %s", err)
|
|
}
|
|
|
|
//fetch bans at current time
|
|
bans, err := ctx.GetBansAt(time.Now())
|
|
if err != nil {
|
|
t.Fatalf("failed to get bans : %s", err)
|
|
}
|
|
|
|
if len(bans) != 0 {
|
|
t.Fatalf("should not have bans, got %d bans", len(bans))
|
|
}
|
|
//get bans in the past
|
|
bans, err = ctx.GetBansAt(time.Now().Add(-24 * 31 * time.Hour))
|
|
if err != nil {
|
|
t.Fatalf("failed to get bans : %s", err)
|
|
}
|
|
|
|
if len(bans) != 1 {
|
|
t.Fatalf("should had 1 ban, got %d bans", len(bans))
|
|
}
|
|
if !reflect.DeepEqual(bans, []map[string]string{map[string]string{
|
|
"source": "local",
|
|
"until": "-720h0m0s",
|
|
"reason": "old_overflow",
|
|
"iptext": "1.2.3.4",
|
|
"cn": "",
|
|
"events_count": "0",
|
|
"action": "ban",
|
|
"as": " ",
|
|
"bancount": "0",
|
|
"scenario": "old_overflow",
|
|
}}) {
|
|
t.Fatalf("unexpected results")
|
|
}
|
|
|
|
}
|
|
|
|
func TestWriteBanApplicationOnly(t *testing.T) {
|
|
validCfg := map[string]string{
|
|
"type": "sqlite",
|
|
"db_path": "./test.db",
|
|
"debug": "false",
|
|
"flush": "true",
|
|
}
|
|
ctx, err := NewDatabase(validCfg)
|
|
if err != nil || ctx == nil {
|
|
t.Fatalf("failed to create simple sqlite")
|
|
}
|
|
|
|
if err := ctx.DeleteAll(); err != nil {
|
|
t.Fatalf("failed to flush existing bans")
|
|
}
|
|
|
|
freshRecordsCount := 12
|
|
|
|
for i := 0; i < freshRecordsCount; i++ {
|
|
//this one expires in the future
|
|
OldSignal := genSignalOccurence(fmt.Sprintf("2.2.2.%d", i))
|
|
|
|
OldSignal.BanApplications[0].Until = time.Now().Add(1 * time.Hour)
|
|
if err = ctx.WriteBanApplication(OldSignal.BanApplications[0]); err != nil {
|
|
t.Fatalf("Failed to insert old signal : %s", err)
|
|
}
|
|
}
|
|
|
|
bans, err := ctx.GetBansAt(time.Now())
|
|
if err != nil {
|
|
t.Fatalf("%s", err)
|
|
}
|
|
if len(bans) != freshRecordsCount {
|
|
t.Fatalf("expected %d, got %d", freshRecordsCount, len(bans))
|
|
}
|
|
}
|