package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/elastic/go-libaudit/v2"
	"github.com/elastic/go-libaudit/v2/auparse"
)

var (
	fs          = flag.NewFlagSet("audit", flag.ExitOnError)
	diag        = fs.String("diag", "", "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")
	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+")
)

type process struct {
	cmdline  string
	rootfs   string
	children []int
}

var pids map[int]*process     //古希腊掌管进程的神
var containers map[string]int // 古希腊掌管容器的神
var wg sync.WaitGroup         // 掌管协程

func main() {
	// 检查用户身份，并添加auditd规则，监听所有syscall
	if os.Geteuid() != 0 {
		fmt.Printf("Err: Please run me as root, %d!\n", os.Getegid())
		return
	}
	syscall := [5]string{"fork", "vfork", "execve", "exit", "exit_group"}
	var auditCmd *exec.Cmd
	auditCmd = exec.Command("auditctl", "-D") // 清空所有规则
	auditCmd.Run()
	// 设置监听规则
	for i := 0; i < 5; i++ {
		auditCmd = exec.Command("auditctl", "-a", "exit,always", "-F", "arch=b64", "-S", syscall[i])
		auditCmd.Run()
	}

	// 查找pid
	containerdPid, err := getPid()
	if err != nil {
		fmt.Printf("Error finding containerd: %v\n", err)
		return
	}

	pids = make(map[int]*process)
	containers = make(map[string]int)
	pids[containerdPid] = &process{cmdline: "/usr/bin/cmdline", rootfs: "/", children: make([]int, 0)}

	// 开始运行，解析命令行参数后监听
	if err := fs.Parse(os.Args[1:]); err != nil {
		log.Fatal(err)
	}

	if err := read(); err != nil {
		log.Fatalf("error: %v", err)
	}
}

func getPid() (int, error) {
	// 指定要搜索的关键词
	keyword := "/usr/bin/containerd"

	// 获取/proc目录下的所有子目录
	procDir, err := filepath.Glob("/proc/*")
	if err != nil {
		return 0, err
	}

	// 遍历子目录，查找包含关键词的进程
	for _, dir := range procDir {
		pid, err := strconv.Atoi(filepath.Base(dir))
		if err != nil {
			continue // 跳过非PID的目录
		}

		// 检查进程是否包含关键词
		if containsKeyword(pid, keyword) {
			return pid, nil
		}
	}
	err = fmt.Errorf("Error: no containerd process found.\n")
	return 0, err
}

func containsKeyword(pid int, keyword string) bool {
	// 构造完整的进程命令路径
	cmdPath := fmt.Sprintf("/proc/%d/cmdline", pid)

	// 打开文件
	file, err := os.Open(cmdPath)
	if err != nil {
		return false
	}
	defer file.Close()

	// 读取文件内容
	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, keyword) {
			return true
		}
	}
	return false
}

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
	}

	log.Println("starting netlink client")

	var err error
	var client *libaudit.AuditClient
	if *receiveOnly {
		client, err = libaudit.NewMulticastAuditClient(diagWriter)
		if err != nil {
			return fmt.Errorf("failed to create receive-only audit client: %w", err)
		}
		defer client.Close()
	} else {
		client, err = libaudit.NewAuditClient(diagWriter)
		if err != nil {
			return fmt.Errorf("failed to create audit client: %w", err)
		}
		defer client.Close()

		status, err := client.GetStatus()
		if err != nil {
			return fmt.Errorf("failed to get audit status: %w", err)
		}
		log.Printf("received audit status=%+v", status)

		if status.Enabled == 0 {
			log.Println("enabling auditing in the kernel")
			if err = client.SetEnabled(true, libaudit.WaitForReply); err != nil {
				return fmt.Errorf("failed to set enabled=true: %w", err)
			}
		}

		if status.RateLimit != uint32(*rate) {
			log.Printf("setting rate limit in kernel to %v", *rate)
			if err = client.SetRateLimit(uint32(*rate), libaudit.NoWait); err != nil {
				return fmt.Errorf("failed to set rate limit to unlimited: %w", err)
			}
		}

		if status.BacklogLimit != uint32(*backlog) {
			log.Printf("setting backlog limit in kernel to %v", *backlog)
			if err = client.SetBacklogLimit(uint32(*backlog), libaudit.NoWait); err != nil {
				return fmt.Errorf("failed to set backlog limit: %w", err)
			}
		}

		if status.Enabled != 2 && *immutable {
			log.Printf("setting kernel settings as immutable")
			if err = client.SetImmutable(libaudit.NoWait); err != nil {
				return fmt.Errorf("failed to set kernel as immutable: %w", err)
			}
		}

		log.Printf("sending message to kernel registering our PID (%v) as the audit daemon", os.Getpid())
		if err = client.SetPID(libaudit.NoWait); err != nil {
			return fmt.Errorf("failed to set audit PID: %w", err)
		}
	}

	return receive(client)
}

func receive(r *libaudit.AuditClient) error {
	for {
		rawEvent, err := r.Receive(false)
		if err != nil {
			return fmt.Errorf("receive failed: %w", err)
		}

		// Messages from 1300-2999 are valid audit messages.
		if rawEvent.Type < auparse.AUDIT_USER_AUTH ||
			rawEvent.Type > auparse.AUDIT_LAST_USER_MSG2 {
			continue
		}

		fmt.Printf("type=%v msg=%s\n", rawEvent.Type, rawEvent.Data)
	}
}