Merge pull request #97 from zema1/feat-config-file

feat: support multiple pusher && config file

CONFIG.md

@@ -0,0 +1,43 @@
+# 配置文件
+
+Watchvuln 从 `v2.0.0` 版本开始支持从文件加载配置, 使用时需要使用 `-c` 参数指定配置文件路径, 如:
+
+```
+./watchvuln -c /path/to/config.yaml
+./watchvuln -c /path/to/config.json
+```
+
+同时为了简化开发和减低理解成本，我们约定，**如果指定了配置文件，那么命令行指定的任何参数将不再生效**
+
+## 文件格式
+
+支持 `yaml` 和 `json` 两种格式的配置文件，这两个格式本质上是互通的，你可以选择自己喜欢的后缀。
+一般，你只需将 `config.example.yaml` 的内容改一下即可，一个最简单的配置大概如下:
+
+```yaml
+db_conn: sqlite3://vuln_v3.sqlite3
+sources: [ "avd","nox","oscs","threatbook","seebug","struts2","kev" ]
+interval: 30m
+pusher:
+  - type: dingding
+    access_token: "xxxx"
+    sign_secret: "yyyy"
+```
+
+聪明的你一定发现了，配置文件里的字段和命令行参数是一一对应的，这里就不再赘述了。
+
+实际上，配置文件的出现主要是为了解决多推送的问题，比如你有两个钉钉群需要推送，那么可以写成这样
+
+```yaml
+db_conn: sqlite3://vuln_v3.sqlite3
+sources: [ "avd","nox","oscs","threatbook","seebug","struts2","kev" ]
+interval: 30m
+pusher:
+  - type: dingding
+    access_token: "xxxx"
+    sign_secret: "yyyy"
+
+  - type: dingding
+    access_token: "pppp"
+    sign_secret: "qqqq"
+```

README.md

@@ -332,6 +332,10 @@ $ ./watchvuln --dt DINGDING_ACCESS_TOKEN --ds DINGDING_SECRET --wk WECHATWORK_KE
 
 </details>
 
+## 配置文件
+
+进入查看详情 [使用配置文件](CONFIG.md)
+
 ## 数据库连接
 
 默认使用 sqlite3 作为数据库，数据库文件为 `vuln_v3.sqlite3`，如果需要使用其他数据库，可以通过 `--db`

config.example.yaml

@@ -0,0 +1,43 @@
+db_conn: sqlite3://vuln_v3.sqlite3
+sources: [ "avd","nox","oscs","threatbook","seebug","struts2","kev" ]
+interval: 30m
+enable_cve_filter: true
+no_github_search: false
+no_start_message: false
+diff_mode: false
+white_keywords: [ ]
+black_keywords: [ ]
+
+pusher:
+  - type: dingding
+    access_token: ""
+    sign_secret: ""
+
+  - type: lark
+    access_token: ""
+    sign_secret: ""
+
+  - type: wechatwork
+    key: ""
+
+  - type: lanxin
+    domain: ""
+    access_token: ""
+    sign_secret: ""
+
+  - type: bark
+    url: ""
+
+  - type: serverchan
+    key: ""
+
+  - type: pushplus
+    token: ""
+
+  - type: telegram
+    bot_token: ""
+    chat_ids: ""
+
+  - type: webhook
+    url: ""
+

ctrl/config.go

@@ -1,25 +1,31 @@
 package ctrl
 
 import (
+	"encoding/json"
+	"entgo.io/ent/dialect"
 	"fmt"
+	"github.com/kataras/golog"
+	"github.com/zema1/watchvuln/push"
 	"net/url"
+	"os"
 	"time"
-
-	"entgo.io/ent/dialect"
 )
 
 type WatchVulnAppConfig struct {
-	DBConn          string        `yaml:"db_conn" json:"db_conn"`
-	Sources         []string      `yaml:"sources" json:"sources"`
-	Interval        time.Duration `yaml:"interval" json:"interval"`
-	EnableCVEFilter bool          `yaml:"enable_cve_filter" json:"enable_cve_filter"`
-	NoGithubSearch  bool          `yaml:"no_github_search" json:"no_github_search"`
-	NoStartMessage  bool          `yaml:"no_start_message" json:"no_start_message"`
-	NoFilter        bool          `yaml:"no_filter" json:"no_filter"`
-	DiffMode        bool          `yaml:"diff_mode" json:"diff_mode"`
-	Version         string        `yaml:"version" json:"version"`
-	WhiteKeywords   []string      `yaml:"white_keywords" json:"white_keywords"`
-	BlackKeywords   []string      `yaml:"black_keywords" json:"black_keywords"`
+	DBConn          string              `yaml:"db_conn" json:"db_conn"`
+	Sources         []string            `yaml:"sources" json:"sources"`
+	Interval        string              `yaml:"interval" json:"interval"`
+	EnableCVEFilter *bool               `yaml:"enable_cve_filter" json:"enable_cve_filter"`
+	NoGithubSearch  *bool               `yaml:"no_github_search" json:"no_github_search"`
+	NoStartMessage  *bool               `yaml:"no_start_message" json:"no_start_message"`
+	DiffMode        *bool               `yaml:"diff_mode" json:"diff_mode"`
+	WhiteKeywords   []string            `yaml:"white_keywords" json:"white_keywords"`
+	BlackKeywords   []string            `yaml:"black_keywords" json:"black_keywords"`
+	Pusher          []map[string]string `yaml:"pusher" json:"pusher"`
+
+	NoFilter       bool          `yaml:"-" json:"-"`
+	Version        string        `yaml:"-" json:"-"`
+	IntervalParsed time.Duration `json:"-" yaml:"-"`
 }
 
 const dbExample = `
@@ -28,6 +34,25 @@ mysql://user:pass@host:port/dbname
 postgres://user:pass@host:port/dbname
 `
 
+func (c *WatchVulnAppConfig) Init() {
+	if c.EnableCVEFilter == nil {
+		t := true
+		c.EnableCVEFilter = &t
+	}
+	if c.NoGithubSearch == nil {
+		c.NoGithubSearch = new(bool)
+	}
+	if c.NoStartMessage == nil {
+		c.NoStartMessage = new(bool)
+	}
+	if c.DiffMode == nil {
+		c.DiffMode = new(bool)
+	}
+	if c.Interval == "" {
+		c.Interval = "1h"
+	}
+}
+
 func (c *WatchVulnAppConfig) DBConnForEnt() (string, string, error) {
 	u, err := url.Parse(c.DBConn)
 	if err != nil {
@@ -64,3 +89,92 @@ func (c *WatchVulnAppConfig) DBConnForEnt() (string, string, error) {
 		return "", "", fmt.Errorf("unsupported db_conn: %s, expected:%s", c.DBConn, dbExample)
 	}
 }
+
+func (c *WatchVulnAppConfig) GetPusher() (push.TextPusher, push.RawPusher, error) {
+	var textPusher []push.TextPusher
+	var rawPusher []push.RawPusher
+
+	for _, config := range c.Pusher {
+		pushType := config["type"]
+
+		switch pushType {
+		case push.TypeDingDing:
+			dingConfig := unmarshal[push.DingDingConfig](config)
+			if dingConfig.AccessToken == "" || dingConfig.SignSecret == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewDingDing(&dingConfig))
+		case push.TypeLark:
+			larkConfig := unmarshal[push.LarkConfig](config)
+			if larkConfig.SignSecret == "" || larkConfig.AccessToken == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewLark(&larkConfig))
+		case push.TypeWechatWork:
+			wechatConfig := unmarshal[push.WechatWorkConfig](config)
+			if wechatConfig.Key == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewWechatWork(&wechatConfig))
+		case push.TypeWebhook:
+			webhookConfig := unmarshal[push.WebhookConfig](config)
+			if webhookConfig.URL == "" {
+				continue
+			}
+			rawPusher = append(rawPusher, push.NewWebhook(&webhookConfig))
+		case push.TypeLanxin:
+			lanxinConfig := unmarshal[push.LanxinConfig](config)
+			if lanxinConfig.Domain == "" || lanxinConfig.AccessToken == "" || lanxinConfig.SignSecret == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewLanxin(&lanxinConfig))
+		case push.TypeBark:
+			barkConfig := unmarshal[push.BarkConfig](config)
+			if barkConfig.URL == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewBark(&barkConfig))
+		case push.TypeServerChan:
+			serverchanConfig := unmarshal[push.ServerChanConfig](config)
+			if serverchanConfig.Key == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewServerChan(&serverchanConfig))
+		case push.TypePushPlus:
+			pushplusConfig := unmarshal[push.PushPlusConfig](config)
+			if pushplusConfig.Token == "" {
+				continue
+			}
+			textPusher = append(textPusher, push.NewPushPlus(&pushplusConfig))
+		case push.TypeTelegram:
+			telegramConfig := unmarshal[push.TelegramConfig](config)
+			if telegramConfig.BotToken == "" || telegramConfig.ChatIDs == "" {
+				continue
+			}
+			tgPusher, err := push.NewTelegram(&telegramConfig)
+			if err != nil {
+				return nil, nil, fmt.Errorf("init telegram error %w", err)
+			}
+			textPusher = append(textPusher, tgPusher)
+		default:
+			return nil, nil, fmt.Errorf("unsupported push type: %s", pushType)
+		}
+		golog.Infof("add pusher: %s", pushType)
+	}
+	if len(textPusher) == 0 && len(rawPusher) == 0 {
+		msg := `
+you must setup at least one pusher, eg: 
+use dingding: %s --dt DINGDING_ACCESS_TOKEN --ds DINGDING_SECRET
+use wechat:   %s --wk WECHATWORK_KEY
+use webhook:  %s --webhook WEBHOOK_URL`
+		return nil, nil, fmt.Errorf(msg, os.Args[0], os.Args[0], os.Args[0])
+	}
+	return push.MultiTextPusher(textPusher...), push.MultiRawPusher(rawPusher...), nil
+}
+
+func unmarshal[T any](config map[string]string) T {
+	data, _ := json.Marshal(config)
+	var res T
+	_ = json.Unmarshal(data, &res)
+	return res
+}

ctrl/ctrl.go

@@ -51,11 +51,16 @@ type WatchVulnApp struct {
 	prs          []*github.PullRequest
 }
 
-func NewApp(config *WatchVulnAppConfig, textPusher push.TextPusher, rawPusher push.RawPusher) (*WatchVulnApp, error) {
+func NewApp(config *WatchVulnAppConfig) (*WatchVulnApp, error) {
+	config.Init()
 	drvName, connStr, err := config.DBConnForEnt()
 	if err != nil {
 		return nil, err
 	}
+	textPusher, rawPusher, err := config.GetPusher()
+	if err != nil {
+		return nil, err
+	}
 	drv, err := entSql.Open(drvName, connStr)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed opening connection to db")
@@ -113,7 +118,7 @@ func NewApp(config *WatchVulnAppConfig, textPusher push.TextPusher, rawPusher pu
 }
 
 func (w *WatchVulnApp) Run(ctx context.Context) error {
-	if w.config.DiffMode {
+	if *w.config.DiffMode {
 		w.log.Info("running in diff mode, skip init vuln database")
 		w.collectAndPush(ctx)
 		w.log.Info("diff finished")
@@ -128,7 +133,7 @@ func (w *WatchVulnApp) Run(ctx context.Context) error {
 		return err
 	}
 	w.log.Infof("system init finished, local database has %d vulns", localCount)
-	if !w.config.NoStartMessage {
+	if !*w.config.NoStartMessage {
 		providers := make([]*grab.Provider, 0, 10)
 		failed := make([]*grab.Provider, 0, 10)
 		for _, p := range w.grabbers {
@@ -140,7 +145,7 @@ func (w *WatchVulnApp) Run(ctx context.Context) error {
 		msg := &push.InitialMessage{
 			Version:        w.config.Version,
 			VulnCount:      localCount,
-			Interval:       w.config.Interval.String(),
+			Interval:       w.config.IntervalParsed.String(),
 			Provider:       providers,
 			FailedProvider: failed,
 		}
@@ -165,11 +170,11 @@ func (w *WatchVulnApp) Run(ctx context.Context) error {
 		time.Sleep(time.Second)
 	}()
 
-	ticker := time.NewTicker(w.config.Interval)
+	ticker := time.NewTicker(w.config.IntervalParsed)
 	defer ticker.Stop()
 	for {
 		w.prs = nil
-		w.log.Infof("next checking at %s\n", time.Now().Add(w.config.Interval).Format("2006-01-02 15:04:05"))
+		w.log.Infof("next checking at %s\n", time.Now().Add(w.config.IntervalParsed).Format("2006-01-02 15:04:05"))
 
 		select {
 		case <-ctx.Done():
@@ -203,7 +208,7 @@ func (w *WatchVulnApp) collectAndPush(ctx context.Context) {
 				w.log.Infof("%s has been pushed, skipped", v)
 				continue
 			}
-			if v.CVE != "" && w.config.EnableCVEFilter {
+			if v.CVE != "" && *w.config.EnableCVEFilter {
 				// 同一个 cve 已经有其它源推送过了
 				others, err := w.db.VulnInformation.Query().
 					Where(vulninformation.And(vulninformation.Cve(v.CVE), vulninformation.Pushed(true))).All(ctx)
@@ -256,7 +261,7 @@ func (w *WatchVulnApp) collectAndPush(ctx context.Context) {
 			}
 
 			// find cve pr in nuclei repo
-			if v.CVE != "" && !w.config.NoGithubSearch {
+			if v.CVE != "" && !*w.config.NoGithubSearch {
 				links, err := w.FindGithubPoc(ctx, v.CVE)
 				if err != nil {
 					w.log.Warn(err)

ctrl/ctrl_test.go

@@ -13,15 +13,11 @@ func TestGithubSearch(t *testing.T) {
 	assert := require.New(t)
 
 	app, err := NewApp(&WatchVulnAppConfig{
-		DBConn:          "sqlite3://vuln_v3.sqlite3",
-		Sources:         nil,
-		Interval:        30,
-		EnableCVEFilter: false,
-		NoGithubSearch:  false,
-		NoStartMessage:  false,
-		NoFilter:        false,
-		Version:         "",
-	}, nil, nil)
+		DBConn:   "sqlite3://vuln_v3.sqlite3",
+		Sources:  nil,
+		Interval: "30h",
+		Version:  "",
+	})
 	assert.Nil(err)
 	links, err := app.FindGithubPoc(context.Background(), "CVE-2023-37582")
 	assert.Nil(err)

go.mod

@@ -18,12 +18,12 @@ require (
 	github.com/kataras/golog v0.1.8
 	github.com/larksuite/oapi-sdk-go/v2 v2.0.11
 	github.com/pkg/errors v0.9.1
-	github.com/rayepeng/serverchan v1.0.0
 	github.com/stretchr/testify v1.8.4
 	github.com/urfave/cli/v2 v2.26.0
 	github.com/vimsucks/wxwork-bot-go v0.0.0-20221213061339-fcbcd88ede1c
 	golang.org/x/net v0.24.0
 	golang.org/x/sync v0.7.0
+	gopkg.in/yaml.v3 v3.0.1
 	modernc.org/sqlite v1.28.0
 )
 
@@ -77,7 +77,6 @@ require (
 	golang.org/x/tools v0.20.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 	lukechampine.com/uint128 v1.3.0 // indirect
 	modernc.org/cc/v3 v3.41.0 // indirect
 	modernc.org/ccgo/v3 v3.16.15 // indirect

go.sum

@@ -138,8 +138,6 @@ github.com/quic-go/qpack v0.4.0 h1:Cr9BXA1sQS2SmDUWjSofMPNKmvF6IiIfDRmgU0w1ZCo=
 github.com/quic-go/qpack v0.4.0/go.mod h1:UZVnYIfi5GRk+zI9UMaCPsmZ2xKJP7XBUvVyT1Knj9A=
 github.com/quic-go/quic-go v0.42.0 h1:uSfdap0eveIl8KXnipv9K7nlwZ5IqLlYOpJ58u5utpM=
 github.com/quic-go/quic-go v0.42.0/go.mod h1:132kz4kL3F9vxhW3CtQJLDVwcFe5wdWeJXXijhsO57M=
-github.com/rayepeng/serverchan v1.0.0 h1:mLrHOwXaqpjkq9GlzsCS5cQM3+yxPdAmIR4VT6mrU7U=
-github.com/rayepeng/serverchan v1.0.0/go.mod h1:eSz+5xrtgpDom0mYpOEaR5MQSvF7VbRGxNoD2qSUwNk=
 github.com/refraction-networking/utls v1.6.0 h1:X5vQMqVx7dY7ehxxqkFER/W6DSjy8TMqSItXm8hRDYQ=
 github.com/refraction-networking/utls v1.6.0/go.mod h1:kHJ6R9DFFA0WsRgBM35iiDku4O7AqPR6y79iuzW7b10=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=