develope v0.4.0

api/routers/v1/plugin.go → api/routers/v1/plugin/plugin.go

@@ -1,6 +1,7 @@
 package v1
 
 import (
+	"bufio"
 	"github.com/astaxie/beego/validation"
 	"github.com/gin-gonic/gin"
 	"github.com/jweny/pocassist/api/msg"
@@ -10,6 +11,13 @@ import (
 	"github.com/jweny/pocassist/poc/rule"
 	"github.com/unknwon/com"
 	"gorm.io/datatypes"
+	"log"
+)
+
+const (
+	TargetUrl = "url"
+	TargetUrlFile = "file"
+	TargetUrlRaw = "raw"
 )
 
 type PluginSerializer struct {
@@ -23,13 +31,21 @@ type PluginSerializer struct {
 	Description   	int            `gorm:"column:description" json:"description"`
 }
 
-type RunPluginSerializer struct {
-	// 运行
-	Target			string		`json:"target"`
+type RunSinglePluginSerializer struct {
+	// 运行单个
+	Target			string		   `json:"target"`
 	Affects         string         `gorm:"column:affects" json:"affects"`
 	JsonPoc         datatypes.JSON `gorm:"column:json_poc" json:"json_poc"`
 }
 
+type RunPluginsSerializer struct {
+	// 批量运行
+	Target			string		`json:"target"`
+	TargetType		string		`json:"target_type"`
+	RunType 		string		`json:"run_type"`
+	VulIdList		[]string	`json:"vul_id_list"`
+}
+
 //获取单个plugin
 func GetPlugin(c *gin.Context) {
 	id := com.StrTo(c.Param("id")).MustInt()
@@ -137,11 +153,12 @@ func UpdatePlugin(c *gin.Context) {
 	valid.Required(plugin.Affects, "Affects").Message("Affects不能为空")
 
 	if ! valid.HasErrors() {
-		if db.ExistPluginByID(plugin.Id){
+		if db.ExistPluginByVulId(plugin.VulId){
+			c.JSON(msg.ErrResp("漏洞编号已存在"))
+			return
+		} else {
 			db.EditPlugin(plugin.Id, plugin)
 			c.JSON(msg.SuccessResp(plugin))
-		} else {
-			c.JSON(msg.ErrResp("record not found"))
 			return
 		}
 	} else {
@@ -172,9 +189,9 @@ func DeletePlugin(c *gin.Context) {
 	}
 }
 
-//运行
+//运行单个plugin 不是从数据库提取数据，表单传数据
 func RunPlugin(c *gin.Context) {
-	run := RunPluginSerializer{}
+	run := RunSinglePluginSerializer{}
 	err := c.BindJSON(&run)
 	if err != nil {
 		c.JSON(msg.ErrResp("参数校验不通过"))
@@ -195,7 +212,7 @@ func RunPlugin(c *gin.Context) {
 				Affects:       run.Affects,
 				JsonPoc:       poc,
 			}
-			item := &rule.ScanItem{Req: oreq, Vul: &currentPlugin}
+			item := &rule.ScanItem{Req: oreq, Plugin: &currentPlugin}
 			result, err := rule.RunPoc(item)
 			if err != nil {
 				c.JSON(msg.ErrResp("规则运行失败：" + err.Error()))
@@ -210,6 +227,62 @@ func RunPlugin(c *gin.Context) {
 	}
 }
 
+//批量运行plugin 从数据库提取数据，表单传数据
+//前端向后端传 "vul_id_list":["poc_db_1","poc_db_2"]
+func RunPlugins(c *gin.Context) {
+	runs := RunPluginsSerializer{}
+	err := c.BindJSON(&runs)
+	if err != nil {
+		c.JSON(msg.ErrResp("参数校验不通过"))
+		return
+	}
+	plugins, err := rule.LoadDbPlugin(runs.RunType, runs.VulIdList)
+
+	switch runs.TargetType {
+	case TargetUrl:
+		url := runs.TargetType
+		oreq, err := util.GenOriginalReq(url)
+		if err != nil {
+			logging.GlobalLogger.Error("[original request gen err ]", err)
+			c.JSON(msg.ErrResp("原始请求生成失败"))
+			return
+		}
+		rule.RunPlugins(oreq, plugins)
+	case TargetUrlFile:
+		//获取文件
+		file, header, err := c.Request.FormFile("file")
+		if err != nil {
+			logging.GlobalLogger.Error("[original request gen err ]", err)
+			c.JSON(msg.ErrResp("url文件上传失败"))
+			return
+		}
+		log.Print(header.Filename)
+		//content, err := ioutil.ReadAll(file)
+		var targets []string
+
+		scanner := bufio.NewScanner(file)
+		for scanner.Scan() {
+			val := scanner.Text()
+			if val == "" {
+				continue
+			}
+			targets = append(targets, val)
+		}
+
+		for _, url := range targets {
+			oreq, err := util.GenOriginalReq(url)
+			if err != nil {
+				logging.GlobalLogger.Error("[original request gen err ]", err)
+			}
+			logging.GlobalLogger.Info("[start check url ]", url)
+			rule.RunPlugins(oreq, plugins)
+		}
+	case TargetUrlRaw:
+		//请求报文
+	}
+}
+
+
 
 
 

cmd/run.go

@@ -8,6 +8,7 @@ import (
 	"github.com/jweny/pocassist/pkg/util"
 	"github.com/jweny/pocassist/poc/rule"
 	"github.com/urfave/cli/v2"
+	"log"
 	"os"
 	"sort"
 )
@@ -18,15 +19,13 @@ var (
 	rawFile		string
 	loadPoc		string
 	condition	string
-	debug		bool
-	dbname		string
 )
 
 func InitAll() {
 	// config 必须最先加载
 	conf2.Setup()
-	logging.Setup(debug)
-	db.Setup(dbname)
+	logging.Setup()
+	db.Setup()
 	routers.Setup()
 	util.Setup()
 	rule.Setup()
@@ -37,21 +36,6 @@ func RunApp() {
 	app.Name = "pocassist"
 	app.Usage = "New POC Framework Without Writing Code"
 	app.Version = "0.3.0"
-	// 全局flag
-	app.Flags = []cli.Flag{
-		&cli.BoolFlag{
-			Name: "debug",
-			Aliases: []string{"d"},
-			Destination: &debug,
-			Value: false,
-			Usage: "enable debug log"},
-		&cli.StringFlag{
-			Name: "database",
-			Aliases: []string{"b"},
-			Destination: &dbname,
-			Value: "sqlite",
-			Usage: "kind of database, default: sqlite"},
-	}
 
 	// 子命令
 	app.Commands = []*cli.Command{
@@ -64,7 +48,7 @@ func RunApp() {
 
 	err := app.Run(os.Args)
 	if err != nil {
-		logging.GlobalLogger.Error("[app run err ]", err)
+		log.Fatalf("cli.RunApp err: %v", err)
 		return
 	}
 }

pkg/conf/config.go

@@ -1,6 +1,8 @@
 package conf
 
 import (
+	"bytes"
+	"github.com/fsnotify/fsnotify"
 	"github.com/spf13/viper"
 	"log"
 	"os"
@@ -64,24 +66,59 @@ var GlobalConfig *Config
 
 // 加载配置
 func Setup() {
-	// 加载config
-	var err error
 	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
 	if err != nil {
-		log.Fatalf("config.Setup, fail to get current path: %v", err)
+		log.Fatalf("conf.Setup, fail to get current path: %v", err)
 	}
+	// 配置文件路径 当前文件夹 + config.yaml
 	configFile := path.Join(dir, "config.yaml")
-	viper.SetConfigFile(configFile)
+
+	// 检测配置文件是否存在
+	_ , err = os.Lstat(configFile)
+	if err != nil {
+		//	没有，生成默认yaml
+		WriteYamlConfig(configFile)
+	}
+	// watch配置
+	ReadYamlConfig(configFile)
+
+}
+
+func ReadYamlConfig(configFile string) {
+	// 加载config
 	viper.SetConfigType("yaml")
+	viper.SetConfigFile(configFile)
 
-	err = viper.ReadInConfig()
+	err := viper.ReadInConfig()
 	if err != nil {
-		log.Fatalf("config.Setup, fail to read 'config.yaml': %v", err)
+		log.Fatalf("conf.Setup, fail to read 'config.yaml': %v", err)
 	}
 	err = viper.Unmarshal(&GlobalConfig)
 	if err != nil {
-		log.Fatalf("config.Setup, fail to parse 'config.yaml': %v", err)
+		log.Fatalf("conf.Setup, fail to parse 'config.yaml', check format: %v", err)
+	}
+	err = verifiyConfig()
+	if err != nil {
+		log.Fatalf("conf.Setup, fail to verify 'config.yaml', check format: %v", err)
 	}
+	// watch 监控配置文件变化
+	viper.WatchConfig()
+	viper.OnConfigChange(func(e fsnotify.Event) {
+		// 配置文件发生变更之后会调用的回调函数
+		log.Println("Config file changed:", e.Name)
+	})
 }
 
-
+func WriteYamlConfig(configFile string) {
+	// 生成默认config
+	viper.SetConfigType("yaml")
+	err := viper.ReadConfig(bytes.NewBuffer(defaultYamlByte))
+	if err != nil {
+		log.Fatalf("conf.Setup, fail to read default config bytes: %v", err)
+	}
+	// 写文件
+	err = viper.SafeWriteConfigAs(configFile)
+	if err != nil {
+		log.Fatalf("conf.Setup, fail to write 'config.yaml': %v", err)
+	}
+}

pkg/conf/default.go

@@ -0,0 +1,84 @@
+package conf
+
+import (
+	"encoding/json"
+	"errors"
+)
+
+var defaultYamlByte = []byte(`
+# webserver配置
+serverConfig:
+  # 配置jwt秘钥
+  jwt_secret: "pocassist"
+  # gin的运行模式 "release" 或者 "debug"
+  run_mode: "release"
+  # 运行日志的文件名，日志将保存在二进制所在目录
+  log_name : "debug.log"
+
+# HTTP配置
+httpConfig:
+  # 扫描时使用的代理：格式为 IP:PORT，example: 如 burpsuite，可填写 127.0.0.1:8080
+  proxy: ""
+  # 读取 http 响应超时时间，不建议设置太小，否则可能影响到盲注的判断
+  http_timeout: 10
+  # 建立 tcp 连接的超时时间
+  dail_timeout: 5
+  # udp 超时时间
+  udp_timeout: 5
+  # 每秒最大请求数
+  max_qps: 100
+  # 单个请求最大允许的跳转次数
+  max_redirect: 5
+  headers:
+    # 默认 UA
+    user_agent: "Mozilla/5.0 (Windows NT 10.0; rv:78.0) Gecko/20100101 Firefox/78.0"
+
+# 数据库配置
+dbConfig:
+  # sqlite配置：sqlite数据库文件的路径
+  sqlite : "pocassist.db"
+  # mysql配置
+  mysql:
+    host: "127.0.0.1"
+    password: ""
+    port: "3306"
+    user: "root"
+    database: "pocassist"
+    # 数据库连接超时时间
+    timeout: "3s"
+
+# 插件配置
+pluginsConfig:
+  # 并发量:同时运行的插件数量
+  parallel: 8
+
+# 反连平台配置: 目前使用 ceye.io
+reverse:
+  api_key: ""
+  domain: ""
+`)
+
+var runMode = []string{"debug","release"}
+
+func ArrayToString (array []string) string {
+	str, _ := json.Marshal(array)
+	return string(str)
+}
+
+func StrInArray (str string, array []string) error {
+	for _, element := range array{
+		if str == element{
+			return nil
+		}
+	}
+	return errors.New(str + "must in" + ArrayToString(array))
+}
+
+func verifiyConfig() error {
+	var err error
+	err = StrInArray(GlobalConfig.ServerConfig.RunMode, runMode)
+	if err != nil {
+		return err
+	}
+	return nil
+}