feat: add datasets watcher (#23)

* feat(internal): add cron pkg

* feat: add datasets watcher

* chore: rm unused file

* feat: add const err msg

* fix: add err check for run.cron

go.mod

@@ -5,7 +5,8 @@ go 1.19
 require (
 	github.com/charmbracelet/log v0.2.4
 	github.com/fsnotify/fsnotify v1.6.0
-	github.com/kitabisa/teler-waf v1.2.0-alpha.1
+	github.com/go-co-op/gocron v1.34.0
+	github.com/kitabisa/teler-waf v1.2.0-beta.1
 	github.com/mattn/go-colorable v0.1.13
 )
 
@@ -20,6 +21,7 @@ require (
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/bgentry/go-netrc v0.0.0-20140422174119-9fd32a8b3d3d // indirect
+	github.com/bitfield/script v0.22.0 // indirect
 	github.com/charmbracelet/lipgloss v0.8.0 // indirect
 	github.com/daniel-hutao/spinlock v0.1.0 // indirect
 	github.com/dwisiswant0/clientip v0.3.0 // indirect
@@ -32,14 +34,16 @@ require (
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/s2a-go v0.1.4 // indirect
-	github.com/google/uuid v1.3.0 // indirect
+	github.com/google/uuid v1.3.1 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.2.5 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.0 // indirect
 	github.com/gorilla/css v1.0.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-getter v1.7.2 // indirect
 	github.com/hashicorp/go-safetemp v1.0.0 // indirect
 	github.com/hashicorp/go-version v1.6.0 // indirect
+	github.com/itchyny/gojq v0.12.12 // indirect
+	github.com/itchyny/timefmt-go v0.1.5 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/klauspost/compress v1.16.7 // indirect
 	github.com/leodido/go-urn v1.2.4 // indirect
@@ -57,6 +61,7 @@ require (
 	github.com/projectdiscovery/mapcidr v1.1.2 // indirect
 	github.com/projectdiscovery/utils v0.0.41 // indirect
 	github.com/rivo/uniseg v0.4.4 // indirect
+	github.com/robfig/cron/v3 v3.0.1 // indirect
 	github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
 	github.com/samber/lo v1.38.1 // indirect
 	github.com/scorpionknifes/go-pcre v0.0.0-20210805092536-77486363b797 // indirect
@@ -66,8 +71,9 @@ require (
 	github.com/valyala/fastjson v1.6.4 // indirect
 	github.com/valyala/fasttemplate v1.2.2 // indirect
 	go.opencensus.io v0.24.0 // indirect
+	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	go.uber.org/zap v1.25.0 // indirect
+	go.uber.org/zap v1.26.0 // indirect
 	golang.org/x/crypto v0.12.0 // indirect
 	golang.org/x/exp v0.0.0-20230626212559-97b1e661b5df // indirect
 	golang.org/x/net v0.14.0 // indirect
@@ -83,4 +89,5 @@ require (
 	google.golang.org/grpc v1.56.2 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	mvdan.cc/sh/v3 v3.6.0 // indirect
 )

internal/cron/cron.go

@@ -0,0 +1,26 @@
+package cron
+
+import (
+	"time"
+
+	"github.com/go-co-op/gocron"
+)
+
+type Cron struct {
+	*gocron.Scheduler
+	*gocron.Job
+}
+
+func New() (*Cron, error) {
+	c := new(Cron)
+
+	tz, err := time.LoadLocation("Local")
+	if err != nil {
+		return c, err
+	}
+
+	c.Scheduler = gocron.NewScheduler(tz)
+	c.Job, err = c.Scheduler.Every(1).Day().At("00:00").WaitForSchedule().Do(task)
+
+	return c, err
+}

internal/cron/task.go

@@ -0,0 +1,19 @@
+package cron
+
+import "github.com/kitabisa/teler-waf/threat"
+
+var task = func() error {
+	updated, err := threat.IsUpdated()
+	if err != nil {
+		return err
+	}
+
+	if !updated {
+		err = threat.Get()
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

internal/runner/const.go

@@ -0,0 +1,5 @@
+package runner
+
+const (
+	errSomething = "Something went wrong"
+)

internal/runner/runner.go

@@ -14,16 +14,21 @@ import (
 	"github.com/charmbracelet/log"
 	"github.com/fsnotify/fsnotify"
 	"github.com/kitabisa/teler-proxy/common"
+	"github.com/kitabisa/teler-proxy/internal/cron"
 	"github.com/kitabisa/teler-proxy/pkg/tunnel"
+	"github.com/kitabisa/teler-waf"
+	"github.com/kitabisa/teler-waf/threat"
 )
 
 type Runner struct {
 	*common.Options
-	*fsnotify.Watcher
+	*cron.Cron
 	*http.Server
 
 	shuttingDown bool
 	shutdownLock sync.Mutex
+	telerOpts    teler.Options
+	watcher
 }
 
 func New(opt *common.Options) error {
@@ -35,23 +40,17 @@ func New(opt *common.Options) error {
 	run := &Runner{Options: opt}
 
 	if opt.Config.Path != "" {
-		watcher, err := fsnotify.NewWatcher()
+		w, err := fsnotify.NewWatcher()
 		if err != nil {
 			return err
 		}
 
-		if err := watcher.Add(opt.Config.Path); err != nil {
+		if err := w.Add(opt.Config.Path); err != nil {
 			return err
 		}
 
-		run.Watcher = watcher
-		defer run.Watcher.Close()
-
-		go func() {
-			if err := run.watch(); err != nil {
-				opt.Logger.Fatal("Something went wrong", "err", err)
-			}
-		}()
+		defer w.Close()
+		run.watcher.config = w
 	}
 
 	dest := buildDest(opt.Destination)
@@ -70,11 +69,42 @@ func New(opt *common.Options) error {
 		Handler:  tun,
 		ErrorLog: logger,
 	}
+
 	run.Server = server
+	run.telerOpts = tun.Options
+
+	if run.shouldCron() && run.Cron == nil {
+		w, err := fsnotify.NewWatcher()
+		if err != nil {
+			return err
+		}
+
+		ds, err := threat.Location()
+		if err != nil {
+			return err
+		}
+
+		if err := w.Add(ds); err != nil {
+			return err
+		}
+
+		defer w.Close()
+		run.watcher.datasets = w
+
+		if err := run.cron(); err != nil {
+			opt.Logger.Fatal(errSomething, "err", err)
+		}
+	}
+
+	go func() {
+		if err := run.watch(); err != nil {
+			opt.Logger.Fatal(errSomething, "err", err)
+		}
+	}()
 
 	go func() {
 		if err := run.start(); err != nil {
-			opt.Logger.Fatal("Something went wrong", "err", err)
+			opt.Logger.Fatal(errSomething, "err", err)
 		}
 	}()
 
@@ -154,13 +184,32 @@ func (r *Runner) notify(sigCh chan os.Signal) error {
 func (r *Runner) watch() error {
 	for {
 		select {
-		case event := <-r.Watcher.Events:
-			if event.Op == 2 {
+		case event := <-r.watcher.config.Events:
+			if event.Op.Has(fsnotify.Write) {
 				r.Options.Logger.Warn("Configuration file has changed", "conf", r.Options.Config.Path)
 				return r.restart()
 			}
-		case err := <-r.Watcher.Errors:
+		case event := <-r.watcher.datasets.Events:
+			if event.Op.Has(fsnotify.Write) || event.Op.Has(fsnotify.Remove) {
+				r.Options.Logger.Warn("Threat datasets has updated", "event", event.Op)
+				return r.restart()
+			}
+		case err := <-r.watcher.config.Errors:
+			return err
+		case err := <-r.watcher.datasets.Errors:
 			return err
 		}
 	}
 }
+
+func (r *Runner) cron() error {
+	c, err := cron.New()
+	if err != nil {
+		return err
+	}
+
+	r.Cron = c
+	c.Scheduler.StartAsync()
+
+	return nil
+}

internal/runner/utils.go

@@ -53,3 +53,17 @@ func isReachable(inputURL string, timeout time.Duration) bool {
 
 	return true
 }
+
+func (r *Runner) shouldCron() bool {
+	if r.Options.Config.Path == "" {
+		return false
+	}
+
+	opt := r.telerOpts
+
+	if !opt.InMemory && !opt.NoUpdateCheck {
+		return true
+	}
+
+	return false
+}

internal/runner/watcher.go

@@ -0,0 +1,7 @@
+package runner
+
+import "github.com/fsnotify/fsnotify"
+
+type watcher struct {
+	config, datasets *fsnotify.Watcher
+}

pkg/tunnel/tunnel.go

@@ -13,8 +13,10 @@ import (
 )
 
 type Tunnel struct {
-	*teler.Teler
 	*httputil.ReverseProxy
+	*teler.Teler
+
+	Options teler.Options
 }
 
 func NewTunnel(port int, dest, cfgPath, optFormat string) (*Tunnel, error) {
@@ -48,6 +50,7 @@ func NewTunnel(port int, dest, cfgPath, optFormat string) (*Tunnel, error) {
 			return nil, err
 		}
 
+		tun.Options = opt
 		tun.Teler = teler.New(opt)
 	} else {
 		tun.Teler = teler.New()