package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"netlink"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

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

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+")
)

func main() {
	// 检查用户身份，并添加auditd规则，监听所有syscall
	if os.Geteuid() != 0 {
		fmt.Printf("Err: Please run me as root, %d!\n", os.Getegid())
		return
	}

	// 所有的系统调用号与名称的关系
	err := figureOutSyscalls()
	if err != nil {
		fmt.Printf("Error figuring out syscall numbers: %v\n", err)
	}

	exec.Command("auditctl", "-D").Run()
	exec.Command("auditctl", "-b", "1000000000").Run()
	exec.Command("auditctl", "--reset-lost").Run()

	var auditCmd *exec.Cmd

	pidSyscall := []string{"execve"}
	// 设置监听规则
	for i := 0; i < len(pidSyscall); i++ {
		auditCmd = exec.Command("auditctl", "-a", "exit,always", "-F", "arch=b64", "-S", pidSyscall[i])
		auditCmd.Run()
	}

	// 监听文件的消息
	fileSyscall := []string{"open", "write", "close"}
	// fileSyscall := []string{"open", "write", "creat", "unlink", "opendir", "mkdir", "rmdir", "chmod", "fchmod", "chown", "fchown", "lchown", "flock"}
	for i := 0; i < len(fileSyscall); i++ {
		auditCmd = exec.Command("auditctl", "-a", "exit,always", "-F", "arch=b64", "-S", fileSyscall[i])
		auditCmd.Run()
	}

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

	// 创世之神，1号进程
	// 1号进程还是不要在进程树上直接出现了，不然它的小儿子们都会出现
	// /usr/bin/containerd，也就是我们最关注的进程

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

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

func coroutine(client *libaudit.AuditClient) {
	// 各协程至此开始
	rawChan = make(chan interface{})
	cookedChan = make(chan Event)

	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)
}

func procWatch() error {
	ns, err := netlink.NewNetlinkSocket(syscall.NETLINK_CONNECTOR, 12345)
	if err != nil {
		fmt.Printf("Error creating socket: %v\n", err)
		return err
	}
	defer ns.Close()
	for {
		res, err := ns.Receive()
		if err != nil {
			fmt.Printf("Error recv: %v\n", err)
			continue
		}
		for i := 0; i < len(res); i++ {
			procEvent := netlink.ParseProcEvent(res[i].Data)
			switch procEvent.What {
			case netlink.PROC_EVENT_FORK:
				data := procEvent.Data.(netlink.ProcEventFork)
				cooked := Event{
					tag:        NEWPID,
					timestamp:  time.Now(),
					pid:        int(data.ChildPid),
					tgid:       int(data.ChildTgid),
					ppid:       int(data.ParentPid),
					parentTgid: int(data.ParentTgid),
				}
				checkProc(&cooked)
				cookedChan <- cooked
			case netlink.PROC_EVENT_EXIT:
				data := procEvent.Data.(netlink.ProcEventExit)
				cooked := Event{
					tag:         PIDEXIT,
					timestamp:   time.Now(),
					pid:         int(data.ProcessPid),
					exit_code:   uint64(data.ExitCode),
					exit_signal: int(data.ExitSignal),
				}
				cookedChan <- cooked
			default:
			}
		}
	}
}

func checkProc(pCooked *Event) {
	fileName := fmt.Sprintf("/proc/%d/task/%d/cmdline", pCooked.tgid, pCooked.pid)
	fd, err := os.Open(fileName)
	if err != nil {
		fmt.Printf("Err: %v\n", err)
		return
	}

	scanner := bufio.NewScanner(fd)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		pCooked.argv = append(pCooked.argv, strings.Split(line, "\x00")...)
	}
	pCooked.argc = len(pCooked.argv)
	fd.Close()

	fileName = fmt.Sprintf("/proc/%d/task/%d/cwd", pCooked.tgid, pCooked.pid)
	pCooked.cwd, err = os.Readlink(fileName)
	if err != nil {
		fmt.Printf("Err readlink %s: %v\n", fileName, err)
		pCooked.cwd = ""
	}
}