6 files changed, 86 insertions, 57 deletions

diff --git a/listener/audit.go b/listener/audit.go
index ed48691..148378c 100644
--- a/listener/audit.go
+++ b/listener/audit.go
@@ -13,14 +13,14 @@ func read() error {
         // Write netlink response to a file for further analysis or for writing
         // tests cases.
         var diagWriter io.Writer
-        if *diag != "" {
-                f, err := os.OpenFile(*diag, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o600)
-                if err != nil {
-                        return err
-                }
-                defer f.Close()
-                diagWriter = f
-        }
+        // if *diag != "" {
+        //      f, err := os.OpenFile(*diag, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o664)
+        //      if err != nil {
+        //              return err
+        //      }
+        //      defer f.Close()
+        //      diagWriter = f
+        // }
 
         log.Println("starting netlink client")
 
diff --git a/listener/deal.go b/listener/deal.go
index 70c2827..af65ff8 100644
--- a/listener/deal.go
+++ b/listener/deal.go
@@ -10,17 +10,8 @@ import (
         "go.mongodb.org/mongo-driver/bson"
 )
 
-const (
-        dbName      string = "test"
-        pidColName  string = "pids"
-        fdColName   string = "fds"
-        fileColName string = "files"
-)
-
-var pidCol, fdCol, fileCol mongoClient
-
 func initPidCol() (err error) {
-        // TODO: 这里是否需要补全一下进程信息？
+        // 这里是否需要补全一下进程信息？
         dirs, err := os.ReadDir(fmt.Sprintf("/proc/%d/task", containerdPid))
         if err != nil {
                 return err
@@ -41,6 +32,9 @@ func initPidCol() (err error) {
                         process.Star = true
                 }
                 err = pidCol.InsertOne(process)
+                if err != nil {
+                        return err
+                }
         }
         return nil
 }
@@ -49,27 +43,6 @@ func deal() {
         defer wg.Done()
         var cooked Event
         var ok bool
-        var err error
-
-        if err = pidCol.init(dbName, pidColName); err != nil {
-                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
-                return
-        }
-        if err = initPidCol(); err != nil {
-                fmt.Fprintf(os.Stderr, "Err while initing pidcol: %v\n", err)
-        }
-
-        if err = fdCol.init(dbName, fdColName); err != nil {
-                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
-                return
-        }
-        if err = fileCol.init(dbName, fileColName); err != nil {
-                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
-        }
-
-        defer pidCol.Disconnect()
-        defer fdCol.Disconnect()
-        defer fileCol.Disconnect()
 
         for {
                 cooked, ok = <-cookedChan
diff --git a/listener/global.go b/listener/global.go
index b782284..49d6e94 100644
--- a/listener/global.go
+++ b/listener/global.go
@@ -44,12 +44,6 @@ type Event struct {
         destPath         string
 }
 
-var wg sync.WaitGroup        // 掌管协程
-var rawChan chan interface{} // 从接收到整理的管道
-var cookedChan chan Event    // 整理好的信息的管道
-var syscallTable [500]string //记录一下系统调用
-var containerdPid int
-
 // 插入到数据库的结构
 type Exec struct {
         Timestamp time.Time `bson:"timestamp"`
@@ -84,3 +78,17 @@ type File struct {
         Written        []time.Time `bson:"written"`
         CloseTimestamp time.Time   `bson:"close_timestamp"`
 }
+
+const (
+        dbName      string = "test"
+        pidColName  string = "pids"
+        fdColName   string = "fds"
+        fileColName string = "files"
+)
+
+var wg sync.WaitGroup                  // 掌管协程
+var rawChan chan interface{}           // 从接收到整理的管道
+var cookedChan chan Event              // 整理好的信息的管道
+var syscallTable [500]string           //记录一下系统调用
+var containerdPid int                  // 容器守护进程进程号
+var pidCol, fdCol, fileCol mongoClient // 数据库集合
diff --git a/listener/godo.go b/listener/godo.go
index 8d82231..0e1dc73 100644
--- a/listener/godo.go
+++ b/listener/godo.go
@@ -18,14 +18,15 @@ import (
 
 var (
         fs          = flag.NewFlagSet("audit", flag.ExitOnError)
-        diag        = fs.String("diag", "", "dump raw information from kernel to file")
+        diag        = fs.String("diag", "godo.log", "dump raw information from kernel to file")
         rate        = fs.Uint("rate", 0, "rate limit in kernel (default 0, no rate limit)")
-        backlog     = fs.Uint("backlog", 8192, "backlog limit")
+        backlog     = fs.Uint("backlog", 1<<30, "backlog limit")
         immutable   = fs.Bool("immutable", false, "make kernel audit settings immutable (requires reboot to undo)")
         receiveOnly = fs.Bool("ro", false, "receive only using multicast, requires kernel 3.16+")
+        mongoURI    = fs.String("mongo", "localhost:27017", "mongo database uri")
 )
 
-const bufferPages = 100
+const bufferPages = 1000
 
 func main() {
         // 检查用户身份，并添加auditd规则，监听所有syscall
@@ -41,7 +42,6 @@ func main() {
         }
 
         exec.Command("auditctl", "-D").Run()
-        exec.Command("auditctl", "-b", "1000000000").Run()
         exec.Command("auditctl", "--reset-lost").Run()
 
         var auditCmd *exec.Cmd
@@ -78,24 +78,45 @@ func main() {
         }
 }
 
-func coroutine(client *libaudit.AuditClient) {
+func coroutine(client *libaudit.AuditClient) error {
         // 各协程至此开始
         bufferSize := bufferPages * syscall.Getpagesize()
         rawChan = make(chan interface{}, bufferSize)
         cookedChan = make(chan Event, bufferSize)
 
+        var err error
+        if err = pidCol.init(dbName, pidColName); err != nil {
+                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
+                return err
+        }
+        if err = initPidCol(); err != nil {
+                fmt.Fprintf(os.Stderr, "Err while initing pidcol: %v\n", err)
+        }
+
+        if err = fdCol.init(dbName, fdColName); err != nil {
+                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
+                return err
+        }
+        if err = fileCol.init(dbName, fileColName); err != nil {
+                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
+        }
+
+        defer pidCol.Disconnect()
+        defer fdCol.Disconnect()
+        defer fileCol.Disconnect()
+
+        wg.Add(1)
+        go deal()
         wg.Add(1)
         go procWatch()
-
         wg.Add(1)
         go receive(client)
         wg.Add(1)
         go orgnaze()
-        wg.Add(1)
-        go deal()
 
         wg.Wait()
         time.Sleep(2 * time.Second)
+        return nil
 }
 
 func procWatch() error {
diff --git a/listener/mongo.go b/listener/mongo.go
index a51350e..36c471c 100644
--- a/listener/mongo.go
+++ b/listener/mongo.go
@@ -31,18 +31,27 @@ func (mc *mongoClient) init(dbName, colName string) error {
 
 func (mc *mongoClient) Connect(dbName, colName string) error {
         var err error
-        mc.client, err = mongo.NewClient(options.Client().ApplyURI("mongodb://localhost:27017"))
+        // 设置连接MongoDB的参数
+        clientOptions := options.Client().ApplyURI("mongodb://" + *mongoURI)
 
+        // 创建一个带有超时的上下文
+        ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+        defer cancel() // 确保在函数退出时取消上下文
+
+        // 使用带超时的上下文连接到MongoDB
+        mc.client, err = mongo.Connect(ctx, clientOptions)
         if err != nil {
                 return err
         }
 
-        ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
-        err = mc.client.Connect(ctx)
+        // 尝试ping数据库以检查连接是否成功
+        err = mc.client.Ping(ctx, nil)
         if err != nil {
                 return err
         }
 
+        fmt.Println("Connected to MongoDB!")
+
         mc.col = mc.client.Database(dbName).Collection(colName)
         mc.dbName = dbName
         mc.colName = colName
diff --git a/listener/organize.go b/listener/organize.go
index 0c05eb4..cf6dad3 100644
--- a/listener/organize.go
+++ b/listener/organize.go
@@ -2,6 +2,7 @@ package main
 
 import (
         "fmt"
+        "io"
         "os"
         "regexp"
         "strconv"
@@ -41,13 +42,30 @@ func orgnaze() {
         var raw interface{}
         var rawEvent libaudit.RawAuditMessage
 
+        var diagWriter io.Writer
+        var f *os.File
+        var err error
+        var fileName string
+        if *diag != "" {
+                fileName = *diag
+        } else {
+                fileName = "godo.log"
+        }
+
+        f, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o664)
+        if err != nil {
+                f = nil
+        }
+        defer f.Close()
+        diagWriter = f
+
         for {
                 raw, ok = <-rawChan
                 if !ok {
                         break
                 }
                 rawEvent = raw.(libaudit.RawAuditMessage)
-                // fmt.Printf("type=%v msg=%s\n", rawEvent.Type, rawEvent.Data)
+                fmt.Fprintf(diagWriter, "type=%v msg=%s\n", rawEvent.Type, rawEvent.Data)
 
                 switch rawEvent.Type {
                 case auparse.AUDIT_SYSCALL: