chore: optmize packet capturing

.gitignore

@@ -15,6 +15,7 @@
 go.work
 
 *.db
+*.db-journal
 *.mmdb
 dist
 /config.toml

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"log"
 	"os"
 	"os/signal"
 	"syscall"
@@ -15,18 +16,11 @@ func main() {
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
 	<-sig
+	log.Println("Shutting down...")
 	util.DB.Close()
 }
 
 func init() {
-	if viper.GetDuration("capture_interval") > 0 {
-		go func() {
-			for {
-				util.Capture()
-				time.Sleep(viper.GetDuration("capture_interval"))
-			}
-		}()
-	}
 	if viper.GetDuration("check_interval") > 0 {
 		go func() {
 			for {

tic/netvigil.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
 
 	"github.com/saurlax/netvigil/util"
@@ -28,21 +29,21 @@ func (t *Netvigil) Check(ips []string) []*util.Threat {
 		IPs:   ips,
 	})
 	if err != nil {
-		fmt.Println("[Netvigil] Failed to marshal request:", err)
+		log.Println("[Netvigil] Failed to marshal request:", err)
 		return threats
 	}
 
 	resp, err := http.Post(fmt.Sprintf("%s/api/check", t.Server), "application/json", bytes.NewBuffer(requestBody))
 	if err != nil {
-		fmt.Println("[Netvigil] Failed to request:", err)
+		log.Println("[Netvigil] Failed to request:", err)
 		return threats
 	}
 	defer resp.Body.Close()
 
 	var res NetvigilResponse
 	err = json.NewDecoder(resp.Body).Decode(&res)
 	if err != nil {
-		fmt.Println("[Netvigil] Failed to decode response:", err)
+		log.Println("[Netvigil] Failed to decode response:", err)
 		return threats
 	}
 

tic/threatbook.go

@@ -2,7 +2,7 @@ package tic
 
 import (
 	"encoding/json"
-	"fmt"
+	"log"
 	"net"
 	"net/http"
 	"net/url"
@@ -41,19 +41,19 @@ func (t *Threatbook) Check(ips []string) []*util.Threat {
 		"resource": resource,
 	})
 	if err != nil {
-		fmt.Println("[Threatbook] Failed to request:", err)
+		log.Println("[Threatbook] Failed to request:", err)
 		return threats
 	}
 	defer resp.Body.Close()
 
 	var res ThreatbookResponse
 	err = json.NewDecoder(resp.Body).Decode(&res)
 	if err != nil {
-		fmt.Println("[Threatbook] Failed to decode response:", err)
+		log.Println("[Threatbook] Failed to decode response:", err)
 		return threats
 	}
 	if res.ResponseCode != 0 {
-		fmt.Printf("[Threatbook] Abnormal response (%v): %v\n", res.ResponseCode, res.VerBoseMsg)
+		log.Printf("[Threatbook] Abnormal response (%v): %v\n", res.ResponseCode, res.VerBoseMsg)
 	}
 
 	for ip, data := range res.Data {

tic/tic.go

@@ -1,9 +1,10 @@
 package tic
 
 import (
-	"fmt"
+	"log"
 	"net"
 
+	"github.com/google/gopacket/layers"
 	"github.com/saurlax/netvigil/util"
 	"github.com/spf13/viper"
 )
@@ -15,8 +16,8 @@ type TIC interface {
 
 var tics = make([]TIC, 0)
 
-// Create a TIC instance with config
-func Create(m map[string]any) TIC {
+// create a TIC instance with config
+func create(m map[string]any) TIC {
 	switch m["type"] {
 	case "local":
 		blacklist := make([]net.IP, 0)
@@ -41,7 +42,7 @@ func Create(m map[string]any) TIC {
 }
 
 // Check all IPs with all TICs created
-func CheckIPs(ips []string) []*util.Threat {
+func CheckAll(ips []string) []*util.Threat {
 	threats, _ := util.GetThreatsByIPs(ips)
 	ips2check := make([]string, 0)
 Loop:
@@ -71,13 +72,22 @@ func Check() {
 Loop:
 	for {
 		select {
-		case ip := <-util.IPs:
-			ips = append(ips, ip)
+		case packet := <-util.Packets:
+			ipv4Layer := packet.Layer(layers.LayerTypeIPv4)
+			if ipv4Layer != nil {
+				ip := ipv4Layer.(*layers.IPv4)
+				ips = append(ips, ip.DstIP.String())
+			}
+			ipv6Layer := packet.Layer(layers.LayerTypeIPv6)
+			if ipv6Layer != nil {
+				ip := ipv6Layer.(*layers.IPv6)
+				ips = append(ips, ip.DstIP.String())
+			}
 		default:
 			break Loop
 		}
 	}
-	CheckIPs(ips)
+	CheckAll(ips)
 }
 
 func init() {
@@ -87,9 +97,9 @@ func init() {
 		if !ok {
 			break
 		}
-		tic := Create(m)
+		tic := create(m)
 		if tic != nil {
-			fmt.Printf("[TIC] %s created\n", m["type"])
+			log.Printf("[TIC] %s created\n", m["type"])
 			tics = append(tics, tic)
 		}
 	}

util/action.go

@@ -2,6 +2,7 @@ package util
 
 import (
 	"fmt"
+	"log"
 	"os"
 	"os/exec"
 
@@ -35,15 +36,15 @@ func DelFireWall(ip string) {
 }
 
 func suspiciousAction(n Netstat) {
-	AddFireWall(n.RemoteIP)
-	fmt.Printf("\x1B[33mSuspicious threat detected: %s → %s\x1B[0m\n", n.Executable, n.RemoteIP)
-	beeep.Notify("Suspicious threat detected!", fmt.Sprintf("%s → %s", n.Executable, n.RemoteIP), "")
+	AddFireWall(n.DstIP)
+	log.Printf("\x1B[33mSuspicious threat detected: %s → %s\x1B[0m\n", n.Executable, n.DstIP)
+	beeep.Notify("Suspicious threat detected!", fmt.Sprintf("%s → %s", n.Executable, n.DstIP), "")
 }
 
 func maliciousAction(n Netstat) {
-	AddFireWall(n.RemoteIP)
-	fmt.Printf("\x1B[31mMalicious threat detected: %s → %s\x1B[0m\n", n.Executable, n.RemoteIP)
-	beeep.Notify("Malicious threat detected!", fmt.Sprintf("%s → %s", n.Executable, n.RemoteIP), "")
+	AddFireWall(n.DstIP)
+	log.Printf("\x1B[31mMalicious threat detected: %s → %s\x1B[0m\n", n.Executable, n.DstIP)
+	beeep.Notify("Malicious threat detected!", fmt.Sprintf("%s → %s", n.Executable, n.DstIP), "")
 }
 
 func (t Threat) Action(n Netstat) {

util/config.go

@@ -1,6 +1,8 @@
 package util
 
 import (
+	"log"
+
 	"github.com/spf13/viper"
 )
 
@@ -12,6 +14,6 @@ func init() {
 	viper.SetConfigFile("config.toml")
 	err := viper.ReadInConfig()
 	if err != nil {
-		panic(err)
+		log.Panicln("Failed to read config:", err)
 	}
 }

util/db.go

@@ -2,6 +2,7 @@ package util
 
 import (
 	"database/sql"
+	"log"
 
 	"github.com/IncSW/geoip2"
 	_ "github.com/mattn/go-sqlite3"
@@ -16,10 +17,10 @@ var (
 func init() {
 	DB, err = sql.Open("sqlite3", "file:netvigil.db")
 	if err != nil {
-		panic(err)
+		log.Panicln("Failed to open database:", err)
 	}
 	GeoLiteCity, err = geoip2.NewCityReaderFromFile("GeoLite2-City.mmdb")
 	if err != nil {
-		panic(err)
+		log.Panicln("Failed to open GeoLite2-City.mmdb:", err)
 	}
 }