package main

import (
	"bufio"
	"debug/elf"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/spf13/cobra"
	"io"
	"os"
	"strings"
)

type syscallArg struct {
	Name string
	Type string
}

func dump(kernelImage, output string) error {
	f, err := elf.Open(kernelImage)
	if err != nil {
		return err
	}
	defer f.Close()

	elfSyms, err := f.Symbols()
	if err != nil {
		return err
	}

	typs := make(map[string]*elf.Symbol)
	args := make(map[string]*elf.Symbol)

	secDataIdx := -1
	secRodataIdx := -1
	var secData, secRodata *elf.Section
	for idx, s := range f.Sections {
		if s.Name == ".data" {
			secDataIdx = idx
			secData = s
		} else if s.Name == ".rodata" {
			secRodataIdx = idx
			secRodata = s
		}
	}

	if secDataIdx == -1 || secRodataIdx == -1 {
		return errors.New("no .data or .rodata section found")
	}
	// fmt.Printf(".data section %+v\n", secData)
	secDataReader := secData.Open()
	// fmt.Printf(".rodata section %+v\n", secData)
	secRodataReader := secRodata.Open()

	for i := range elfSyms {
		s := &elfSyms[i]
		if strings.HasPrefix(s.Name, "types__") {
			syscall := strings.TrimPrefix(s.Name, "types__")
			typs[syscall] = s
		} else if strings.HasPrefix(s.Name, "args__") {
			syscall := strings.TrimPrefix(s.Name, "args__")
			args[syscall] = s
		}
	}

	ReadString := func(dataAddr uint64) (string, error) {
		if _, err = secDataReader.Seek(int64(dataAddr-secData.Addr), io.SeekStart); err != nil {
			return "", err
		}

		var v uint64
		if err = binary.Read(secDataReader, f.FileHeader.ByteOrder, &v); err != nil {
			return "", err
		}

		if _, err = secRodataReader.Seek(int64(v-secRodata.Addr), io.SeekStart); err != nil {
			return "", err
		}

		brd := bufio.NewReader(secRodataReader)
		var data []byte
		if data, err = brd.ReadBytes(0); err != nil {
			return "", err
		}

		return string(data[:len(data)-1]), nil
	}

	info := make(map[string][]syscallArg, 0)
	for name := range typs {
		typSym := typs[name]
		argSym := args[name]
		if int(typSym.Section) != secDataIdx || int(argSym.Section) != secDataIdx {
			panic("symbols not in .data")
		}
		if (typSym.Size != argSym.Size) || (typSym.Size%8 != 0) {
			panic("invalid symbol size")
		}

		syscallArgs := make([]syscallArg, 0)
		n := typSym.Size / 8
		for i := uint64(0); i < n; i++ {
			var typStr, argStr string

			typPtr := typSym.Value + (i * 8)
			if typStr, err = ReadString(typPtr); err != nil {
				return err
			}

			argPtr := argSym.Value + (i * 8)
			if argStr, err = ReadString(argPtr); err != nil {
				return err
			}

			syscallArgs = append(syscallArgs, syscallArg{
				Name: argStr,
				Type: typStr,
			})
		}
		info[name] = syscallArgs
	}

	data, err := json.MarshalIndent(info, "", "   ")
	if err != nil {
		return err
	}

	out := os.Stdout
	if output != "" {
		if out, err = os.Create(output); err != nil {
			return err
		}
		defer out.Close()
	}

	_, err = out.Write(data)
	return err
}

func main() {
	var output string

	cmd := &cobra.Command{
		Use:  "dump_syscalls [options] kernel_image",
		Long: "Dump syscall prototypes from a kernel image, and debug symbols are required in image file",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return errors.New("kernel image is required")
			}

			return dump(args[0], output)
		},
		DisableFlagsInUseLine: true,
		SilenceUsage:          true,
		SilenceErrors:         true,
	}

	cmd.PersistentFlags().StringVarP(&output, "output", "o", "", "specify output file to store syscall prototypes")
	if err := cmd.Execute(); err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "Dump syscall prototypes failed, %v\n", err)
		os.Exit(1)
	}
}