package main

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

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

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 Event struct {
	timestamp time.Time
	pid, ppid int
	syscall   int
	argc      int
	args      []string
	cwd       string
}

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

var pids map[int]*process     //古希腊掌管进程的神
var containers map[string]int // 古希腊掌管容器的神
var wg sync.WaitGroup         // 掌管协程
var rawChan chan interface{}  // 从接收到整理的管道，这里不是原始数据类型，下文解释
var cookedChan chan Event     // 整理好的信息的管道

var syscallTable [500]string //记录一下系统调用

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

	syscall := [6]string{"fork", "vfork", "clone", "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)

	// 创世之神，1号进程
	pids[1] = &process{rootfs: "/", children: make([]int, 0)}
	pids[1].children = append(pids[1].children, containerdPid)
	// /usr/bin/containerd，也就是我们最关注的进程
	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 figureOutSyscalls() error {
	NRRegex := regexp.MustCompile(`#define __NR_(.*?) (\d+)$`)
	file, err := os.Open("/usr/include/asm/unistd_64.h")
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		if NRRegex.MatchString(line) {
			match := NRRegex.FindStringSubmatch(line)
			num, err := strconv.Atoi(match[2])
			if err != nil {
				return err
			}
			syscallTable[num] = match[1]
		}
	}
	return nil
}

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 getTimeFromStr(timeStr string) (time.Time, error) {
	timestampFloat, err := strconv.ParseFloat(timeStr, 64)
	if err != nil {
		return time.Unix(0, 0), err
	}
	secs := int64(timestampFloat)
	nsecs := int64((timestampFloat - float64(secs)) * 1e9)

	// 只精确到毫秒就够了
	t := time.Unix(secs, nsecs).Truncate(time.Millisecond)
	return t, nil
}

func hexToAscii(hexString string) string {
	bytes := []byte{}
	for i := 0; i < len(hexString); i += 2 {
		hexPair := hexString[i : i+2]
		// 将十六进制数转换为十进制数
		decimal, err := strconv.ParseInt(hexPair, 16, 8)
		if err != nil {
			return "Invalid hex string"
		}
		char := byte(decimal)
		bytes = append(bytes, char)
	}

	asciiString := strings.ReplaceAll(string(bytes), "\000", " ")

	return asciiString
}

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)
	rawChan = make(chan interface{})
	cookedChan = make(chan Event)
	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 receive(r *libaudit.AuditClient) error {
	defer wg.Done()
	defer close(rawChan)
	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
		}

		rawEventMessage := deepcopy.Copy(*rawEvent)
		rawChan <- rawEventMessage
	}
}

func orgnaze() {
	defer wg.Done()
	defer close(cookedChan)
	// 接收信息
	var raw interface{}
	var ok bool
	var rawEvent libaudit.RawAuditMessage
	// 事件信息
	var eventId, argc int
	var err [6]error
	var event, cooked Event
	// 为每个事务id存储其信息，事务id在操作系统运行期间是唯一的
	eventTable := make(map[int]*Event)
	// 要用的正则匹配列表
	syscallRegex := regexp.MustCompile(`audit\((\d+\.\d+):(\d+)\).*?syscall=(\d+).*?ppid=(\d+) pid=(\d+).*?$`)
	execveRegex := regexp.MustCompile(`audit\(\d+\.\d+:(\d+)\): argc=(\d+)`)
	argsRegex := regexp.MustCompile(`a\d+=("(.*?)"|([0-9a-fA-F]+))`)
	cwdRegex := regexp.MustCompile(`audit\(\d+\.\d+:(\d+)\):  cwd="(.*?)"`)
	proctitleRegex := regexp.MustCompile(`audit\(\d+\.\d+:(\d+)\): proctitle=("(.*?)"|([0-9a-fA-F]+))$`)
	eoeRegex := regexp.MustCompile(`audit\(\d+\.\d+:(\d+)\)`)
	for {
		raw, ok = <-rawChan
		if !ok {
			break
		}
		rawEvent = raw.(libaudit.RawAuditMessage)

		// type Event struct {
		// 	timestamp time.Time
		// 	pid, ppid int
		// 	syscall   int
		// 	argc      int
		// 	args      []string
		// 	cwd       string
		// }
		switch rawEvent.Type {
		case auparse.AUDIT_SYSCALL:
			if syscallRegex.Match(rawEvent.Data) {
				match := syscallRegex.FindSubmatch(rawEvent.Data)
				event.timestamp, err[0] = getTimeFromStr(string(match[1]))
				eventId, err[1] = strconv.Atoi(string(match[2]))
				event.syscall, err[2] = strconv.Atoi(string(match[3]))
				event.ppid, err[3] = strconv.Atoi(string(match[4]))
				event.pid, err[4] = strconv.Atoi(string(match[5]))
				eventTable[eventId] = &Event{
					timestamp: event.timestamp,
					syscall:   event.syscall,
					ppid:      event.ppid,
					pid:       event.pid,
					argc:      0,
					args:      make([]string, 0),
					cwd:       "",
				}
			}
		case auparse.AUDIT_EXECVE:
			if execveRegex.Match(rawEvent.Data) {
				match := execveRegex.FindSubmatch(rawEvent.Data)
				eventId, err[0] = strconv.Atoi(string(match[1]))
				argc, err[1] = strconv.Atoi(string(match[2]))
				if err[0] == nil && err[1] == nil && argsRegex.Match(rawEvent.Data) {
					match := argsRegex.FindAllSubmatch(rawEvent.Data, -1)
					for i := 0; i < argc; i++ {
						if len(match[i][2]) == 0 {
							// 代表着匹配到的是十六进制数
							str := hexToAscii(string(match[i][3]))
							eventTable[eventId].args = append(eventTable[eventId].args, str)
							fmt.Printf("Origin: \"%s\", Res: \"%s\"\n", match[i][3], str)
						} else {
							eventTable[eventId].args = append(eventTable[eventId].args, string(match[i][2]))
						}
					}
					eventTable[eventId].argc = argc
				}
			}
		// case auparse.AUDIT_PATH:
		case auparse.AUDIT_CWD:
			if cwdRegex.Match(rawEvent.Data) {
				match := cwdRegex.FindSubmatch(rawEvent.Data)
				eventId, err[0] = strconv.Atoi(string(match[1]))
				eventTable[eventId].cwd = string(match[2])
			}
		case auparse.AUDIT_PROCTITLE:
			if proctitleRegex.Match(rawEvent.Data) {
				var cmdline string
				var pEvent *Event
				match := proctitleRegex.FindSubmatch(rawEvent.Data)
				eventId, err[0] = strconv.Atoi(string(match[1]))
				pEvent = eventTable[eventId]
				if pEvent.argc == 0 {
					// 只有等于0，才证明没经过EXECVE提取参数，才允许使用PROCTITLE提取参数
					if match[3] == nil {
						// PROCTITLE写的是十六进制，转换为字符串
						cmdline = hexToAscii(string(match[4]))
					} else {
						cmdline = string(match[3])
					}
					pEvent.args = strings.Split(cmdline, " ")
					pEvent.argc = len(eventTable[eventId].args)
				}
				// 当读到proctitle的时候，而且是个新进程最好检查一下cwd，如果还为空，找proc
				if pEvent.cwd == "" && (pEvent.syscall == 57 || pEvent.syscall == 58 || pEvent.syscall == 59) {
					cwdFilePath := fmt.Sprintf("/proc/%d/cwd", pEvent.pid)
					pEvent.cwd, err[1] = os.Readlink(cwdFilePath)
					if err[1] != nil {
						pEvent.cwd = ""
						break
					}
				}
			}
		case auparse.AUDIT_EOE:
			if eoeRegex.Match(rawEvent.Data) {
				match := eoeRegex.FindSubmatch(rawEvent.Data)
				eventId, err[0] = strconv.Atoi(string(match[1]))
				// TODO: 事件整理完毕，即刻发出，是否合理呢？
				cooked = *eventTable[eventId] // 应当采用深拷贝吗？有待实验
				cookedChan <- cooked
				delete(eventTable, eventId) //发出之后就从信息表扔掉，死人别占地
			}
		default:
			// TODO: 这里也需要做防护
		}
	}
}

func deal() {
	defer wg.Done()
	var cooked Event
	var ok bool
	for {
		cooked, ok = <-cookedChan
		if !ok {
			break
		}
		// type Event struct {
		// 	timestamp time.Time
		// 	pid, ppid int
		// 	syscall   int
		// 	argc      int
		// 	args      []string
		// 	cwd       string
		// }
		fmt.Printf("recv: %v syscall=%d, ppid=%d, pid=%d, cwd=\"%s\", argc=%d, ", cooked.timestamp, cooked.syscall, cooked.ppid, cooked.pid, cooked.cwd, cooked.argc)
		if len(cooked.args) != cooked.argc {
			fmt.Printf("Fuck!\n")
			continue
		}
		for i := 0; i < cooked.argc; i++ {
			fmt.Printf("arg[%d]=\"%s\", ", i, cooked.args[i])
		}
		fmt.Printf("\n")
	}
}