package aotinject

/*
#cgo CFLAGS: -I ../inject/include
#cgo amd64 LDFLAGS: ${SRCDIR}/../inject/lib/libhotpatch_amd64.a
#cgo arm64 LDFLAGS: ${SRCDIR}/../inject/lib/libhotpatch_arm64.a
#include "hotpatch.h"
#include <stdlib.h>
*/
import "C"

import (
	"bufio"
	"debug/dwarf"
	"debug/elf"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"syscall"
	"unsafe"
)

func InjectLib(pid int, libPath string) int {
	dll := C.CString(libPath) // 替换为实际的DLL路径
	fmt.Printf("dll: %s\n", libPath)
	rootfs := C.CString("")
	defer C.free(unsafe.Pointer(dll)) // 确保在使用完字符串后释放内存
	result := C.cw_inject_library(C.int(pid), C.int(0), dll, rootfs)
	fmt.Printf("Result JVM: %d\n", result)
	return int(result)
}

// readMemory 用于读取指定地址的内存数据
func readMemory(pid int, address uint64, size uint64) ([]byte, error) {
	memFile := fmt.Sprintf("/proc/%d/mem", pid)
	file, err := os.Open(memFile)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	data := make([]byte, size)
	_, err = file.ReadAt(data, int64(address))
	if err != nil {
		return nil, err
	}

	return data, nil
}

func getFunctionOffsetDBG(libPath, functionName string) (elf.Symbol, error) {
	dwarfFile, err := elf.Open(libPath)
	if err != nil {
		return elf.Symbol{}, fmt.Errorf("failed to open libPath %s: %v", libPath, err)
	}
	dwarfData, err := dwarfFile.DWARF()
	if err != nil {
		return elf.Symbol{}, fmt.Errorf("failed to read DWARF data: %v", err)
	}
	type uprobesDef struct {
		Name       string
		Offset     uint64
		EntAddress uint64
		RetAddress uint64
	}

	listEntry := make(map[dwarf.Offset]uprobesDef)
	SpecListEntry := []dwarf.Entry{}

	entryReader := dwarfData.Reader()

	for {
		entry, err := entryReader.Next()
		if err == io.EOF {
			// We've reached the end of DWARF entries
			break
		}
		if err != nil {
			log.Println("Error reading entry: %v", err)
		}
		if entry == nil {
			log.Println("Warning: a nil entry was returned with no error")
			break
		}
		// fmt.Printf("-----entry address: %v, %v, %v, %v, %v\n", entry, entry.Val(dwarf.AttrName), entry.Tag, entry.Val(dwarf.AttrLowpc), entry.Val(dwarf.AttrHighpc))
		if entry.Tag == dwarf.TagSubprogram {
			// fmt.Printf("entry address: %x, %d\n", entry.Offset, entry.Children)
			funName, ok := entry.Val(dwarf.AttrName).(string)
			if !ok {
				continue
			}
			if functionName == funName {
				entAddress, ok := entry.Val(dwarf.AttrLowpc).(uint64)
				if !ok {
					entAddress = 0
				}
				retAddress, ok := entry.Val(dwarf.AttrHighpc).(uint64)
				if !ok {
					retAddress = 0
				}
				fmt.Printf("Function %s address: %x, %x, %x\n", funName, entAddress, entry.Offset, retAddress)
				uprobes := uprobesDef{
					Name:       funName,
					Offset:     uint64(entry.Offset),
					EntAddress: entAddress,
					RetAddress: retAddress,
				}

				listEntry[entry.Offset] = uprobes
			}

			specAddr, _ := entry.Val(dwarf.AttrSpecification).(dwarf.Offset)
			lowpc := entry.Val(dwarf.AttrLowpc)
			if lowpc != nil && specAddr > 0 && lowpc.(uint64) > 0 {
				SpecListEntry = append(SpecListEntry, *entry)
			}
		}
	}

	for _, v := range SpecListEntry {
		specAddr, _ := v.Val(dwarf.AttrSpecification).(dwarf.Offset)
		// fmt.Printf("SpecListEntrySpecListEntrySpecListEntry Attach Function: %x\n", specAddr)
		_, ok := listEntry[specAddr]
		if ok {
			vv := listEntry[specAddr]
			entAddr := v.Val(dwarf.AttrLowpc)
			if entAddr != nil {
				vv.EntAddress = entAddr.(uint64)
			}
			retAddr := v.Val(dwarf.AttrHighpc)
			if retAddr != nil {
				switch retAddr.(type) {
				case uint64:
					vv.RetAddress = uint64(retAddr.(uint64))
				case int64:
					vv.RetAddress = uint64(retAddr.(int64))
				default:
					fmt.Println("Unknown type")
				}
			}
			listEntry[specAddr] = vv
		}
	}

	for _, v := range listEntry {
		sSize := v.RetAddress
		if v.RetAddress > v.EntAddress {
			sSize = v.RetAddress - v.EntAddress
		}
		symbol := elf.Symbol{
			Name:  v.Name,
			Value: v.EntAddress,
			Size:  180,
		}
		fmt.Printf("Need Attach Function %s address: %x, %x, %d\n", v.Name, v.EntAddress, v.RetAddress, sSize)
		return symbol, nil
	}
	return elf.Symbol{}, fmt.Errorf("function %s not found", functionName)
}

func getFunctionOffset(libPath, functionName string) (elf.Symbol, error) {
	elfFile, err := elf.Open(libPath)
	if err != nil {
		return elf.Symbol{}, fmt.Errorf("failed to open ELF file: %v", err)
	}
	defer elfFile.Close()

	// 优先从 Symbols 中查找函数的地址，如果找不到则从 DynamicSymbols 中查找
	symbols, err := elfFile.Symbols()
	if err != nil {
		fmt.Printf("Can Not find Symbols: %s, %s, %v, can find from DynamicSymbols\n", libPath, functionName, err)
		symbols, err = elfFile.DynamicSymbols()
		if err != nil {
			return elf.Symbol{}, fmt.Errorf("failed to read dynamic symbols: %v", err)
		}
	}

	for _, sym := range symbols {
		if sym.Name == functionName {
			return sym, nil
		}
	}
	return elf.Symbol{}, fmt.Errorf("function %s not found from Symbols", functionName)
}

// getProcessMapsInfo reads the first line of /proc/<pid>/maps file and returns the memory map as a MemoryMap struct
// If a library name is provided, it returns the start address of that library
func getProcessMapsInfo(pid int, libraryName ...string) (*ProcessMapsInfo, error) {
	file, err := os.Open(fmt.Sprintf("/proc/%d/maps", pid))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		addresses := strings.Split(fields[0], "-")
		if len(addresses) != 2 {
			return nil, fmt.Errorf("unexpected format in /proc/%d/maps", pid)
		}

		start, err := strconv.ParseUint(addresses[0], 16, 64)
		if err != nil {
			return nil, err
		}

		end, err := strconv.ParseUint(addresses[1], 16, 64)
		if err != nil {
			return nil, err
		}

		if len(libraryName) == 0 || (len(fields) > 5 && strings.Contains(fields[5], libraryName[0])) {
			return &ProcessMapsInfo{
				Start: start,
				End:   end,
				Path:  fields[5],
			}, nil
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return nil, fmt.Errorf("specified library or process start address not found in /proc/%d/maps", pid)
}

func GetProcessFunctionInfo(pid int, libPath string, functionName string, libraryName ...string) (*ProcessFunctionInfo, error) {
	processMapsInfo, _ := getProcessMapsInfo(pid, libraryName...)

	baseAddress := uint64(processMapsInfo.Start)

	// functionName 是以 @plt 结尾的函数名，则调用 getFunctionOffsetPLT 获取函数的地址，并且去掉 @plt 后缀
	if strings.HasSuffix(functionName, "@plt") {
		functionName = functionName[:len(functionName)-4]
		sendFunctionNew, _ := getFunctionOffsetPLT(pid, libPath, functionName)
		return &ProcessFunctionInfo{
			Name:  functionName,
			Offset: sendFunctionNew.Value,
			Start: baseAddress + sendFunctionNew.Value,
			Size:  sendFunctionNew.Size,
		}, nil
	}

	// 获取函数的偏移量，有限直接从 ELF 文件中获取，如果失败则从 DWARF 文件中获取
	functionSym, err := getFunctionOffset(libPath, functionName)
	if err != nil {
		fmt.Printf("Error getting function offset: %v, now to find dbg %s\n", err, libPath+".dbg")
		functionSym, err = getFunctionOffsetDBG(libPath+".dbg", functionName)
		if err != nil {
			fmt.Printf("Error getting function offset from DWARF file: %v\n", err)
			return nil, err
		}
	}
	fmt.Printf("get ProcessFunctionInfo: 0x%x, 0x%x\n", baseAddress, functionSym.Value)
	funRelAddr := functionSym.Value
	if functionSym.Value < baseAddress {
		funRelAddr += baseAddress
	}
	return &ProcessFunctionInfo{
		Name:   functionSym.Name,
		Offset: functionSym.Value,
		Start:  funRelAddr,
		Size:   functionSym.Size,
	}, nil
}

func AotInject(aotInjector AOTInjector) (int, error) {
	isInject := false

	// lib 注入
	cwLibraryInfo, err := getProcessMapsInfo(aotInjector.PID, aotInjector.CWLibName)
	if err != nil {
		fmt.Printf("Not Find Lib %s: Now to InjectLib: %v\n", aotInjector.CWLibName, err)
		InjectLib(aotInjector.PID, aotInjector.CWLibPath)
		cwLibraryInfo, err = getProcessMapsInfo(aotInjector.PID, aotInjector.CWLibName)
		if err != nil {
			fmt.Printf("InjectLib %s Err: %v\n", aotInjector.CWLibName, err)
			return 0, fmt.Errorf("InjectLib %s Err: %v", aotInjector.CWLibName, err)
		}
	} else {
		isInject = true
	}
	fmt.Printf("find %s Lib Start: %x, End: %x, Path: %s\n", aotInjector.CWLibName, cwLibraryInfo.Start, cwLibraryInfo.End, cwLibraryInfo.Path)
	// 查找 cwFunctionName 函数的偏移量
	cwSym, _ := GetProcessFunctionInfo(aotInjector.PID, cwLibraryInfo.Path, aotInjector.CWFunctionName, aotInjector.CWLibName)
	// 打印 cwFunctionName 函数的偏移量
	fmt.Printf("%s: addr %x, size %d\n", aotInjector.CWFunctionName, cwSym.Start, cwSym.Size)

	if isInject == true {
		fmt.Printf("cwLibrary is already Injected, to find hookOffset\n")
		hookOffset, err := findNopFunctionHookOffset(aotInjector.PID, cwSym)
		if err != nil {
			fmt.Printf("findNopFunctionHookOffset Err: %v Need Inject\n", err)
		} else {
			fmt.Printf("findNopFunctionHookOffset: %d[%d]\n", hookOffset, hookOffset+nopEntryOffset)
			return hookOffset + nopEntryOffset, nil
		}
	}

	// 获取socketWrite0 和 NET_Send 函数的偏移量
	jvmInjectorFunction, err := GetProcessFunctionInfo(aotInjector.PID, aotInjector.InjectLibPath, aotInjector.InjectFunctionName)
	if err != nil {
		fmt.Printf("Error getting function %s offset: %v\n", aotInjector.InjectFunctionName, err)
		return 0, fmt.Errorf("error getting function %s offset: %v", aotInjector.InjectFunctionName, err)
	}
	sendFunction, err := GetProcessFunctionInfo(aotInjector.PID, aotInjector.SendFunctionLibPath, aotInjector.SendFunctionName)
	if err != nil {
		fmt.Printf("Error getting function %s offset: %v\n", aotInjector.SendFunctionName, err)
		return 0, fmt.Errorf("error getting function %s offset: %v", aotInjector.SendFunctionName, err)
	}

	fmt.Printf("find %s, addr: %x, size:%d\n", aotInjector.InjectFunctionName, jvmInjectorFunction.Start, jvmInjectorFunction.Size)
	fmt.Printf("find %s, addr: %x, size:%d\n", aotInjector.SendFunctionName, sendFunction.Start, sendFunction.Size)

	// 读取 socketWrite0 函数的内存数据并解析 NET_Send 在 socketWrite0 中的调用位置
	code, err := readMemory(aotInjector.PID, jvmInjectorFunction.Start, jvmInjectorFunction.Size)
	if err != nil {
		fmt.Printf("Error reading jvmInjectorFunction memory: %v\n", err)
		return 0, fmt.Errorf("error reading jvmInjectorFunction memory: %v", err)
	}
	// 打印 socketWrite0 函数的内存数据
	fmt.Printf("socketWrite0 code: %v\n", code)
	SendFunctionAbsOffset, SendFunctionRelOffset, err := findSendFunctionNameAddr(jvmInjectorFunction.Start, code, aotInjector.SendFunctionName, sendFunction.Start)

	if err != nil {
		fmt.Printf("findSendFunctionNameAddr Err: %v\n", err)
		return 0, fmt.Errorf("findSendFunctionNameAddr Err: %v", err)
	}

	fmt.Printf("find %s : SendFunctionAbsOffset:%x, SendFunctionRelOffset: %d\n", aotInjector.SendFunctionName, SendFunctionAbsOffset, SendFunctionRelOffset)
	// 尝试获取 NET_Send 附近的数据凑够 17 个字节，用于长跳转到 libmylib 中
	instMemStruct, err := findMemForLongJump(jvmInjectorFunction.Start, code, aotInjector.SendFunctionName, SendFunctionAbsOffset)
	if err != nil {
		fmt.Printf("findMemForLongJump Err: %v\n", err)
		return 0, fmt.Errorf("findMemForLongJump Err: %v", err)
	}
	// 打印 instMemStruct 结构体的内容
	fmt.Printf("instMemStruct: %v\n", instMemStruct)

	// 生成长跳指令，跳转到 libmylib 中的 cwFunctionName 函数的+12地址
	jumpCode := generateLongJumpCode(cwSym.Start+nopEntryOffset, len(instMemStruct.CodeArray), true)
	// 使用 parseAndPrintInstructions打印 jumpCode 的内容
	fmt.Printf("jumpCode: %v\n", jumpCode)
	parseAndPrintInstructions(jumpCode)
	// 生成新的指令，用于替换原来的指令，一个参数是原来的指令 instCodeArray，另一个参数是跳转回去的地址，第三个参数是原始 Send 函数的地址，第四个参数是在哪个指令前插入自定义指令
	newCode, hookOffset := generateNewCode(instMemStruct.InstCodeArray, instMemStruct.InstStartAddr+jumpBackAddrOffset, sendFunction.Start, instMemStruct.InsertIndex)
	if hookOffset == 0 {
		fmt.Printf("generateNewCode Err: %v\n", instMemStruct)
		return 0, fmt.Errorf("generateNewCode Err: %v", instMemStruct)
	}
	// 打印 hookIndex
	fmt.Printf("ebpf should hookIndex: %d[%d]\n", hookOffset, nopEntryOffset+hookOffset)
	// 使用 parseAndPrintInstructions 打印 newCode 的内容
	fmt.Printf("newCode: %v\n", newCode)
	parseAndPrintInstructions(newCode)
	// 使用 ptrace attach 目标进程
	fmt.Printf("Attach Process: %d\n", aotInjector.PID)
	if err := syscall.PtraceAttach(aotInjector.PID); err != nil {
		fmt.Printf("PtraceAttach Err: %v\n", err)
		return 0, fmt.Errorf("ptrace attach: %v", err)
	}
	// 等待目标进程停止
	if _, err := syscall.Wait4(aotInjector.PID, nil, 0, nil); err != nil {
		fmt.Printf("wait4: %v", err)
		if err = syscall.PtraceDetach(aotInjector.PID); err != nil {
			fmt.Printf("ptrace DETACH: %v", err)
		}
		return 0, fmt.Errorf("wait4: %v", err)
	}
	fmt.Printf("Ptrace 将newCode使用ptrace写入到目标进程的内存中: %x\n", cwSym.Start+nopEntryOffset)
	// 将newCode使用ptrace写入到目标进程的的空白内存中，地址为 cwSym.Start+12
	_, err = syscall.PtracePokeData(aotInjector.PID, uintptr(cwSym.Start+nopEntryOffset), newCode)
	if err != nil {
		fmt.Printf("PtracePokeData Err: %v\n", err)
	}
	fmt.Printf("Ptrace 将jumpCode使用ptrace写入到目标进程的内存中: %x\n", instMemStruct.InstStartAddr)
	// 将jumpCode使用ptrace写入到目标进程的内存中
	_, err = syscall.PtracePokeData(aotInjector.PID, uintptr(instMemStruct.InstStartAddr), jumpCode)
	if err != nil {
		fmt.Printf("PtracePokeData Err: %v\n", err)
	}
	fmt.Printf("Ptrace DETACH: %d\n", aotInjector.PID)
	// 恢复执行
	if err = syscall.PtraceDetach(aotInjector.PID); err != nil {
		fmt.Printf("ptrace DETACH: %v", err)
	}
	return hookOffset + nopEntryOffset, nil
}

type AOTInjector struct {
	PID                 int
	SendFunctionLibPath string
	SendFunctionName    string
	InjectLibPath       string
	InjectFunctionName  string
	CWLibName           string
	CWFunctionName      string
	CWLibPath           string
}

func main() {
	flag.StringVar(&PID, "p", "", "PID")
	flag.Parse()
	pidStr := PID // 替换为目标进程的 PID
	pid, err := strconv.Atoi(pidStr)
	if err != nil {
		log.Fatalf("Invalid PID: %v", err)
	}

	// NETCoreAOTInject
	// netCoreAotInjector := AOTInjector{
	// 	PID:                 pid,
	// 	SendFunctionLibPath: "/data/NET8/CoreAoT/bin/Debug/net8.0/linux-x64/publish/CoreAoT",
	// 	SendFunctionName:    "send@plt",
	// 	InjectLibPath:       "/data/NET8/CoreAoT/bin/Debug/net8.0/linux-x64/publish/CoreAoT",
	// 	InjectFunctionName:  "SystemNative_Send",
	// 	CWLibName:           "libmylib.so",
	// 	CWFunctionName:      "asmnop",
	// 	CWLibPath:           "/data/NET8/CoreAoT/bin/Debug/net8.0/linux-x64/publish/CoreAoT",
	// }
	// JavaAOTInject
	javaAotInjector := AOTInjector{
		PID:                 pid,
		SendFunctionLibPath: "/data/roger/ebpfdemo/graalve_demo/SimpleHttpServer",
		SendFunctionName:    "NET_Send",
		InjectLibPath:       "/data/roger/ebpfdemo/graalve_demo/SimpleHttpServer",
		InjectFunctionName:  "Java_java_net_SocketOutputStream_socketWrite0",
		CWLibName:           "libmylib.so",
		CWFunctionName:      "asmnop",
		CWLibPath:           "/data/roger/ebpfdemo/mylib/libmylib.so",
	}

	hookOffset, err := AotInject(javaAotInjector)
	if err != nil {
		fmt.Printf("AotInject Err: %v\n", err)
	}
	fmt.Printf("AotInject result hookOffset: %d\n", hookOffset)
}