From d6c6e13796435f9e1e59fec891aa53680748a2d7 Mon Sep 17 00:00:00 2001
From: We-unite <3205135446@qq.com>
Date: Tue, 30 Jul 2024 19:37:48 +0800
Subject: Try to use kernel connector

---
 src/deal.go   | 57 ++++++++++++++++++++++++++++++++++++++++++++-------------
 src/global.go | 48 ++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 86 insertions(+), 19 deletions(-)

(limited to 'src')

diff --git a/src/deal.go b/src/deal.go
index aaac8c5..717344c 100644
--- a/src/deal.go
+++ b/src/deal.go
@@ -55,18 +55,25 @@ func deal() {
 			break
 		}
 
+		// fmt.Printf("%v\n", cooked)
+
 		switch cooked.tag {
 		case NEWPID:
 			dealNewPid(cooked)
 		case EXECVE:
+			check(cooked)
 			dealExecve(cooked)
 		case PIDEXIT:
+			check(cooked)
 			deletePid(cooked)
 		case FILEOPEN:
+			check(cooked)
 			fileOpen(cooked)
 		case FILEWRITE:
+			check(cooked)
 			fileWrite(cooked)
 		case FILECLOSE:
+			check(cooked)
 			fileClose(cooked)
 		}
 	}
@@ -168,19 +175,19 @@ func dealExecve(cooked Event) {
 				},
 			},
 		})
-	} else {
-		// 先fork抵达，插入
-		pidCol.InsertOne(bson.M{
-			"ppid":     cooked.ppid,
-			"pid":      cooked.pid,
-			"children": []bson.M{},
-			"execve": []bson.M{
-				{
-					"timestamp": cooked.timestamp,
-					"execArgs":  cooked.argv,
-				},
-			},
-		})
+		// } else {
+		// 	// 先fork抵达，插入
+		// 	pidCol.InsertOne(bson.M{
+		// 		"ppid":     cooked.ppid,
+		// 		"pid":      cooked.pid,
+		// 		"children": []bson.M{},
+		// 		"execve": []bson.M{
+		// 			{
+		// 				"timestamp": cooked.timestamp,
+		// 				"execArgs":  cooked.argv,
+		// 			},
+		// 		},
+		// 	})
 	}
 	mongoMutex.Unlock()
 }
@@ -255,3 +262,27 @@ func fileWrite(cooked Event) {
 		"close_timestamp": bson.M{"$exists": false},
 	}, bson.M{"$push": bson.M{"written": cooked.timestamp}})
 }
+
+func check(cooked Event) {
+	// 检查进程是否需要记录
+	// 有无父进程在观察中
+	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.ppid})
+	if err != nil || len(docRes) != 1 {
+		return
+	}
+
+	// 自身是否已经记录
+	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.pid})
+	if err != nil {
+		fmt.Printf("Err finding: %v\n", err)
+		return
+	}
+	if len(docRes) == 0 {
+		pidCol.InsertOne(bson.M{
+			"ppid":            cooked.ppid,
+			"pid":             cooked.pid,
+			"children":        []bson.M{},
+			"start_timestamp": cooked.timestamp,
+		})
+	}
+}
diff --git a/src/global.go b/src/global.go
index d1c5c0f..f0f909c 100644
--- a/src/global.go
+++ b/src/global.go
@@ -1,8 +1,11 @@
 package main
 
 import (
+	"fmt"
 	"sync"
 	"time"
+
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 type eventType int
@@ -12,11 +15,19 @@ const (
 	PIDEXIT
 	EXECVE
 	FILEOPEN
-	FILEWRITE
 	FILECLOSE
+	FILEWRITE
 	TYPENUM
 )
 
+func (et eventType) String() string {
+	names := []string{"NEWPID", "PIDEXIT", "EXECVE", "FILEOPEN", "FILECLOSE", "FILEWRITE", "TYPENUM"}
+	if et < NEWPID || et > TYPENUM {
+		return "Unknown"
+	}
+	return names[et]
+}
+
 type Event struct {
 	tag          eventType
 	timestamp    time.Time
@@ -30,12 +41,37 @@ type Event struct {
 	pathName     string
 }
 
-func (et eventType) String() string {
-	names := []string{"NEWPID", "PIDEXIT", "EXECVE", "FILEOPEN", "FILEWRITE", "TYPENUM"}
-	if et < NEWPID || et > TYPENUM {
-		return "Unknown"
+func (event Event) String() string {
+	var res string
+	res = fmt.Sprintf("tag: %v\ntimestamp: %v\nppid: %d\npid: %d\n", event.tag, event.timestamp.Local(), event.ppid, event.pid)
+	res += fmt.Sprintf("syscall: %s\nexit_code: %d\nargs: \n", syscallTable[event.syscall], event.exit_code)
+	for i := 0; i < len(event.argv); i++ {
+		res += fmt.Sprintf("\t\"%s\"\n", event.argv[i])
 	}
-	return names[et]
+	res += "syscallParam: "
+	for i := 0; i < len(event.syscallParam); i++ {
+		res += fmt.Sprintf("\t\"%d\"\n", event.syscallParam[i])
+	}
+	res += "pathName: \"" + event.pathName + "\"\n------\n"
+	return res
+}
+
+type pidExec struct {
+	timestamp time.Time `bson:"timestamp"`
+	execArgs  []string  `bson:"execArgs"`
+}
+
+type pid struct {
+	ID              primitive.ObjectID `bson:"_id,ometempty"`
+	start_timestamp time.Time          `bson:"start_timestamp"`
+	ppid            int                `bson:"ppid"`
+	pid             int                `bson:"pid"`
+	cwd             string             `bson:"cwd"`
+	args            []string           `bson:"args"`
+	execve          []pidExec          `bson:"execve"`
+	children        []int              `bson:"children"`
+	exit_timestamp  time.Time          `bson:"exit_timestamp"`
+	exit_code       uint64             `bson:"exit_code"`
 }
 
 var wg sync.WaitGroup        // 掌管协程
-- 
cgit v1.2.3-70-g09d2