add: ai config

Author: amirhnajafiz

internal/config/config.go

@@ -8,6 +8,7 @@ import (
 	"github.com/ptaas-tool/base-api/internal/config/ftp"
 	"github.com/ptaas-tool/base-api/internal/config/migration"
 	"github.com/ptaas-tool/base-api/internal/config/scanner"
+	"github.com/ptaas-tool/base-api/internal/core/ai"
 	"github.com/ptaas-tool/base-api/internal/storage/sql"
 
 	"github.com/knadh/koanf"
@@ -23,6 +24,7 @@ type Config struct {
 	Migrate migration.Config `koanf:"migrate"`
 	Scanner scanner.Config   `koanf:"scanner"`
 	FTP     ftp.Config       `koanf:"ftp"`
+	AI      ai.Config        `koanf:"ai"`
 }
 
 func Load(path string) Config {

internal/config/default.go

@@ -5,6 +5,7 @@ import (
 	"github.com/ptaas-tool/base-api/internal/config/ftp"
 	"github.com/ptaas-tool/base-api/internal/config/migration"
 	"github.com/ptaas-tool/base-api/internal/config/scanner"
+	"github.com/ptaas-tool/base-api/internal/core/ai"
 	"github.com/ptaas-tool/base-api/internal/storage/sql"
 )
 
@@ -34,5 +35,11 @@ func Default() Config {
 			Secret: "",
 			Access: "",
 		},
+		AI: ai.Config{
+			Enable: true,
+			Method: "random",
+			Limit:  10,
+			Factor: 7,
+		},
 	}
 }

internal/core/ai/ai.go

@@ -2,21 +2,27 @@ package ai
 
 import "math/rand"
 
-type AI struct{}
+type AI struct {
+	Cfg Config
+}
 
 func (a AI) GetAttacks(list, vulnerabilities []string) []string {
 	records := make([]string, 0)
 
-	// logic goes here (now it's random)
-	for _, item := range list {
-		if item == "injection" || item == "payload" {
-			records = append(records, item)
-
-			continue
-		}
+	if !a.Cfg.Enable {
+		return list
+	}
 
-		if rand.Intn(10) > 8 {
-			records = append(records, item)
+	switch a.Cfg.Method {
+	case "svm":
+		break
+	case "nbias":
+		break
+	default:
+		for _, item := range list {
+			if rand.Intn(a.Cfg.Limit) > a.Cfg.Factor {
+				records = append(records, item)
+			}
 		}
 	}
 

internal/core/ai/config.go

@@ -0,0 +1,8 @@
+package ai
+
+type Config struct {
+	Method string `koanf:"method"`
+	Factor int    `koanf:"factor"`
+	Limit  int    `koanf:"limit"`
+	Enable bool   `koanf:"enable"`
+}