package logic

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/elastic/go-libaudit/v2"
	"github.com/elastic/go-libaudit/v2/aucoalesce"
	"github.com/elastic/go-libaudit/v2/auparse"
	"github.com/shirou/gopsutil/v4/process"
)

var (
	// 接收审计事件的管道
	EventChan = make(chan *aucoalesce.Event, 8192)

	// 	记录需要注意的系统调用
	Syscalls = map[string]FileAction{
		"open":      FA_Open,
		"openat":    FA_Open,
		"close":     FA_Close,
		"rename":    FA_Rename,
		"renameat":  FA_Rename,
		"renameat2": FA_Rename,
		"link":      FA_Rename,
		"unlink":    FA_Rename,
		"linkat":    FA_Rename,
		"unlinkat":  FA_Rename,
	}

	SCP = "/usr/bin/scp"

	EventMap     = make(map[string]*EventSet)
	EventMapLock = sync.RWMutex{} // 保护EventMap的锁

	ErrArgNil = errors.New("error arg is nil")
)

type FileAction uint8

const (
	FA_Open FileAction = iota
	FA_Close
	FA_Rename
	FA_Delete
)

func CleanSCP() {
	EventMapLock.Lock()
	defer EventMapLock.Unlock()
	toDelete := make([]string, len(EventMap))
	for k, v := range EventMap {
		alive, _ := v.IsAlive()
		if !alive {
			toDelete = append(toDelete, k)
		}
	}
	for _, v := range toDelete {
		delete(EventMap, v)
	}
}

type EventHandler struct{}

func (h *EventHandler) ReassemblyComplete(msgs []*auparse.AuditMessage) {
	event, err := aucoalesce.CoalesceMessages(msgs)
	if err != nil {
		log.Printf("coalesce messages error: %v \n", err)
	}
	EventChan <- event
}

func (h *EventHandler) EventsLost(count int) {
	log.Printf("event lost: %d \n", count)
}

func StartAuditMonitor(wg *sync.WaitGroup) {
	defer wg.Done()
	cli, err := libaudit.NewMulticastAuditClient(nil)
	if err != nil {
		log.Printf("failed to create audit client: %v \n", err)
		return
	}
	defer cli.Close()

	handler := &EventHandler{}
	rea, err := libaudit.NewReassembler(8192, time.Second*180, handler)
	if err != nil {
		log.Printf("%v", err)
		return
	}
	defer rea.Close()

	// 定期维护事件池
	go func() {
		ticker := time.NewTicker(time.Second * 15)
		defer ticker.Stop()
		for range ticker.C {
			if rea.Maintain() != nil {
				break
			}
		}
	}()

	go FiltAuditMsg()

	for {
		rawMsg, err := cli.Receive(false)
		if err != nil {
			break
		}
		err = rea.Push(rawMsg.Type, slices.Clone(rawMsg.Data))
		if err != nil {
			log.Printf("push audit event error: %v\n", err)
		}
	}
	close(EventChan)
}

// EventSet 事件集，一个事件集记录了同一个进程的审计事件
type EventSet struct {
	Pid, PPID, Uid int32                         // pid、ppid和Uid
	Events         []*aucoalesce.Event           // 记录所有事件
	Fds            map[int64][]*aucoalesce.Event // 记录可以获取文件描述符的事件
	Lock           sync.RWMutex                  // 锁，用于保护Fds和Events的读写
}

// PushEvent 向事件集中插入事件，并判断是否需要触发文件扫描
func (es *EventSet) PushEvent(event *aucoalesce.Event, syscall string) {
	if event == nil {
		return
	}
	switch syscall {
	case "open", "openat":
		fd, err := getExitFd(event)
		if err != nil {
			log.Println("error: ", err)
			return
		}
		if fd == nil {
			log.Println("error: open(at) syscall, but can't get fd")
			return
		}
		es.Lock.Lock()
		el, have := es.Fds[*fd]
		if have && len(el) > 0 {
			// 已经用了这个fd，有问题，要么是没有删除旧fd，要么是有bug
			bb, _ := json.Marshal(el)
			cc, _ := json.Marshal(event)
			log.Println("error: fd repeat, old: ", string(bb), "new: ", string(cc))
			clear(el)
			el = append(el, event)
			es.Events = append(es.Events, event)
			es.Lock.Unlock()
		} else {
			// ok，正常的
			el = make([]*aucoalesce.Event, 0, 16)
			el = append(el, event)
			es.Fds[*fd] = el
			es.Events = append(es.Events, event)
			es.Lock.Unlock()
		}
	case "close":
		// 获取close的第一个参数，那就是fd
		if event.Data == nil {
			return
		}
		fdStr, have := event.Data["a0"]
		if !have {
			log.Println("error: syscall close, but can't find the first arg")
			return
		}
		fd, err := ParseInt(fdStr)
		if err != nil {
			log.Println("error: ", err)
			return
		}
		es.Lock.Lock()
		// 判断fd是否存在，存在就扫描文件
		es.Events = append(es.Events, event)
		el, have := es.Fds[fd]
		if have {
			delete(es.Fds, fd)
			es.Lock.Unlock()
			l := len(el)
			if l == 0 {
				// bug
				log.Println("error: fds slices's length is 0")
				return
			}
			filePath, err := getFilePath(el[l-1])
			if err != nil {
				log.Println("error: ", err)
				return
			}
			if filePath == nil {
				log.Println("error: find file path is nil")
				return
			}
			go es.ScanFile(*filePath)
		} else {
			// 不存在就返回
			es.Events = append(es.Events, event)
			es.Lock.Unlock()
		}
	}
}

func (es *EventSet) ScanFile(path string) {
	log.Printf("user UID(%d) upload file: %s, scanning...\n", es.Uid, path)
	v, _ := ScanFile(path)
	if v {
		log.Printf("user UID(%d) upload file %s containing viruses \n", es.Uid, path)
	} else {
		log.Printf("user UID(%d) upload file %s not find virus \n", es.Uid, path)
	}
}

// NewEventSet 从open、openat系统调用创建事件集
func NewEventSet(event *aucoalesce.Event, _ string) (*EventSet, error) {
	info, err := getBaseInfo(event)
	if err != nil {
		return nil, err
	}
	result := EventSet{
		Events: make([]*aucoalesce.Event, 0, 16),
		Fds:    make(map[int64][]*aucoalesce.Event),
		Lock:   sync.RWMutex{},
	}
	result.Events = append(result.Events, event)
	if info[0] != nil {
		result.Pid = *info[0]
	}
	if info[1] != nil {
		result.PPID = *info[1]
	}
	if info[2] != nil {
		result.Uid = *info[2]
	}
	if event.Data != nil {
		exit, have := event.Data["exit"]
		if !have {
			return nil, errors.New("create file syscall, but not find exit code")
		}
		fd, err := ParseInt(exit)
		if err != nil {
			return nil, err
		}
		el := make([]*aucoalesce.Event, 0, 16)
		el = append(el, event)
		result.Fds[fd] = el
	}
	return &result, nil
}

// CheckAlive 检查进程是否存活
func (es *EventSet) IsAlive() (bool, error) {
	p, err := process.NewProcess(es.Pid)
	if err != nil {
		// 进程已经结束
		return false, nil
	}
	ppid, err := p.Ppid()
	if err != nil {
		return false, nil
	}
	if ppid == es.PPID {
		return true, nil
	}
	return false, nil
}

// PushEvent 向EventMap中放入事件
func PushEvent(syscall string, event *aucoalesce.Event) {
	if event == nil {
		return
	}
	key := genKey(event)
	switch syscall {
	case "open", "openat":
		// 新建文件
		EventMapLock.Lock()
		es, have := EventMap[key]
		if !have {
			ies, err := NewEventSet(event, syscall)
			if err != nil {
				log.Println(err)
				EventMapLock.Unlock()
				return
			}
			es = ies
			EventMap[key] = es
		}
		EventMapLock.Unlock()
	case "close":
		// 关闭文件
		EventMapLock.Lock()
		es, have := EventMap[key]
		if !have {
			EventMapLock.Unlock()
			return
		}
		EventMapLock.Unlock()
		es.PushEvent(event, syscall)
	}
}

// FiltAuditMsg 过滤出有用的事件，存放到map中，并触发处理程序
func FiltAuditMsg() {
	for i := range EventChan {
		if i.Process.Exe != SCP {
			continue
		}
		if i.Category != aucoalesce.EventTypeAuditRule {
			// 不是由审计规则触发的事件跳过
			continue
		}
		if i.Type != auparse.AUDIT_SYSCALL {
			// 不是系统调用，跳过
			continue
		}
		sc, have := i.Data["syscall"]
		if !have {
			continue
		}
		// 检查是不是需要关注的系统调用
		_, have = Syscalls[sc]
		if !have {
			continue
		}
		// 失败
		if i.Result != "success" {
			continue
		}
		// printEvent(i)
		PushEvent(sc, i)
	}
}

func printEvent(e *aucoalesce.Event) {
	if e == nil {
		return
	}
	switch e.Data["syscall"] {
	case "open", "openat":
		fd, _ := getExitFd(e)
		fp, _ := getFilePath(e)
		log.Printf("open(at): pid: %s ,fd: %d, file: %s \n", e.Process.PID, *fd, *fp)
	case "close":
		log.Printf("close: pid: %s ,fd: %s \n", e.Process.PID, e.Data["a0"])
	default:
		log.Println("unknown syscall: ", e.Data["syscall"])
	}
	// bb, err := json.MarshalIndent(e, "", "  ")
	// if err == nil {
	// 	fmt.Println(string(bb))
	// }
}

func ParseInt(str string) (int64, error) {
	if strings.ContainsAny(str, "abcdefABCDEFxX") {
		// 是16进制数
		s := strings.ToLower(str)
		if strings.HasPrefix(s, "0x") {
			return strconv.ParseInt(s, 0, 64)
		} else {
			return strconv.ParseInt(fmt.Sprintf("0x%s", s), 0, 64)
		}
	} else {
		// 不是16进制数
		n, err := strconv.ParseInt(str, 10, 64)
		if err != nil {
			return 0, err
		}
		return n, err
	}
}

// genKey 为事件生成字符串id，格式为 pid-ppid
func genKey(event *aucoalesce.Event) string {
	return fmt.Sprintf("%s-%s", event.Process.PID, event.Process.PPID)
}

// getBaseInfo 获取事件的基本信息，pid、ppid、uid
func getBaseInfo(event *aucoalesce.Event) ([3]*int32, error) {
	result := [3]*int32{nil, nil, nil}
	if event == nil {
		return result, ErrArgNil
	}
	var pid, ppid, uid int32
	if event.User.IDs != nil {
		u, have := event.User.IDs["uid"]
		if have {
			uu, err := ParseInt(u)
			if err == nil {
				uid = int32(uu)
				result[2] = &uid
			}
		}
	}
	p, err := ParseInt(event.Process.PID)
	if err == nil {
		pid = int32(p)
		result[0] = &pid
	}
	p, err = ParseInt(event.Process.PPID)
	if err == nil {
		ppid = int32(p)
		result[0] = &ppid
	}
	return result, nil
}

func getFilePath(event *aucoalesce.Event) (*string, error) {
	if event == nil {
		return nil, ErrArgNil
	}
	if event.File == nil {
		return nil, nil
	}
	return &event.File.Path, nil
}

// getExitFd 获取系统调用的返回值
func getExitFd(event *aucoalesce.Event) (*int64, error) {
	if event == nil {
		return nil, ErrArgNil
	}
	if event.Data != nil {
		exit, have := event.Data["exit"]
		if have {
			e, err := ParseInt(exit)
			if err != nil {
				return nil, err
			}
			return &e, nil
		}
	}
	return nil, nil
}