mdouchement / shigoto

A nextgen crontab.
MIT License
0 stars 0 forks source link

tengo: Enhance shigoto lib #18

Open mdouchement opened 5 months ago

mdouchement commented 5 months ago

Add the following features:

mdouchement commented 5 months ago

Some POC code used to inject host binary stdout & stderr instead of Golang's default ones:

package runner

import (
    "fmt"
    "io"
    "os"
    osexec "os/exec"
    "syscall"

    tengopkg "github.com/d5/tengo/v2"
    "github.com/d5/tengo/v2/stdlib"
    "github.com/mdouchement/ldt/pkg/tengolib"
    "github.com/mdouchement/logger"
    "github.com/pkg/errors"
)

func withfmt(modules *tengopkg.ModuleMap, w io.Writer) {
    m := modules.GetBuiltinModule("fmt")
    // Overwrite useful methods defined in v2.16.1

    m.Attrs["print"] = &tengopkg.UserFunction{
        Name: "print",
        Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
            if len(args) < 1 {
                return nil, tengopkg.ErrWrongNumArguments
            }

            fmt.Fprint(w, tengolib.InterfaceArray(args)...)
            return tengopkg.UndefinedValue, nil
        },
    }
    m.Attrs["printf"] = &tengopkg.UserFunction{
        Name: "printf",
        Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
            if len(args) < 2 {
                return nil, tengopkg.ErrWrongNumArguments
            }

            str, err := tengolib.Format(args...)
            if err != nil {
                return nil, err
            }

            fmt.Fprint(w, str)
            return tengopkg.UndefinedValue, nil
        },
    }
    m.Attrs["println"] = &tengopkg.UserFunction{
        Name: "println",
        Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
            if len(args) < 1 {
                return nil, tengopkg.ErrWrongNumArguments
            }

            fmt.Fprintln(w, tengolib.InterfaceArray(args)...)
            return tengopkg.UndefinedValue, nil
        },
    }
}

func withexec(modules *tengopkg.ModuleMap, dir string, env map[string]string, wout io.Writer, werr io.Writer) {
    m := modules.GetBuiltinModule("os")
    // Overwrite useful methods defined in v2.16.1

    m.Attrs["exec"] = &tengopkg.UserFunction{
        Name: "exec",
        Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
            if len(args) == 0 {
                return nil, tengopkg.ErrWrongNumArguments
            }

            bin, ok := tengopkg.ToString(args[0])
            if !ok {
                return nil, tengopkg.ErrInvalidArgumentType{
                    Name:     "first",
                    Expected: "string(compatible)",
                    Found:    args[0].TypeName(),
                }
            }
            bin, err := osexec.LookPath(bin)
            if err != nil {
                return tengolib.WrapError(err), nil
            }

            arguments := make([]string, 0, len(args)-1)
            for idx, arg := range args[1:] {
                execArg, ok := tengopkg.ToString(arg)
                if !ok {
                    return nil, tengopkg.ErrInvalidArgumentType{
                        Name:     fmt.Sprintf("args[%d]", idx),
                        Expected: "string(compatible)",
                        Found:    args[1+idx].TypeName(),
                    }
                }
                arguments = append(arguments, execArg)
            }

            cmd := &osexec.Cmd{
                Path: bin,
                Args: arguments,
                Dir:  dir,
            }

            if wout != nil {
                cmd.Stdout = wout
            }
            if werr != nil {
                cmd.Stderr = werr
            }

            for k, v := range env {
                cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
            }

            fmt.Printf("%#v\n", cmd)
            return makeOSExecCommand(cmd), nil
        },
    }
}

func makeOSExecCommand(cmd *osexec.Cmd) *tengopkg.ImmutableMap {
    return &tengopkg.ImmutableMap{
        Value: map[string]tengopkg.Object{
            // combined_output() => bytes/error
            "combined_output": &tengopkg.UserFunction{
                Name:  "combined_output",
                Value: stdlib.FuncARYE(cmd.CombinedOutput),
            },
            // output() => bytes/error
            "output": &tengopkg.UserFunction{
                Name:  "output",
                Value: stdlib.FuncARYE(cmd.Output),
            }, //
            // run() => error
            "run": &tengopkg.UserFunction{
                Name:  "run",
                Value: stdlib.FuncARE(cmd.Run),
            }, //
            // start() => error
            "start": &tengopkg.UserFunction{
                Name:  "start",
                Value: stdlib.FuncARE(cmd.Start),
            }, //
            // wait() => error
            "wait": &tengopkg.UserFunction{
                Name:  "wait",
                Value: stdlib.FuncARE(cmd.Wait),
            }, //
            // set_path(path string)
            "set_path": &tengopkg.UserFunction{
                Name: "set_path",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 1 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }
                    s1, ok := tengopkg.ToString(args[0])
                    if !ok {
                        return nil, tengopkg.ErrInvalidArgumentType{
                            Name:     "first",
                            Expected: "string(compatible)",
                            Found:    args[0].TypeName(),
                        }
                    }
                    cmd.Path = s1
                    return tengopkg.UndefinedValue, nil
                },
            },
            // set_dir(dir string)
            "set_dir": &tengopkg.UserFunction{
                Name: "set_dir",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 1 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }
                    s1, ok := tengopkg.ToString(args[0])
                    if !ok {
                        return nil, tengopkg.ErrInvalidArgumentType{
                            Name:     "first",
                            Expected: "string(compatible)",
                            Found:    args[0].TypeName(),
                        }
                    }
                    cmd.Dir = s1
                    return tengopkg.UndefinedValue, nil
                },
            },
            // set_env(env array(string))
            "set_env": &tengopkg.UserFunction{
                Name: "set_env",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 1 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }

                    var env []string
                    var err error
                    switch arg0 := args[0].(type) {
                    case *tengopkg.Array:
                        env, err = tengolib.StringArray(arg0.Value, "first")
                        if err != nil {
                            return nil, err
                        }
                    case *tengopkg.ImmutableArray:
                        env, err = tengolib.StringArray(arg0.Value, "first")
                        if err != nil {
                            return nil, err
                        }
                    default:
                        return nil, tengopkg.ErrInvalidArgumentType{
                            Name:     "first",
                            Expected: "array",
                            Found:    arg0.TypeName(),
                        }
                    }
                    cmd.Env = append(cmd.Env, env...) // Custom modification behavior
                    return tengopkg.UndefinedValue, nil
                },
            },
            // process() => imap(process)
            "process": &tengopkg.UserFunction{
                Name: "process",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 0 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }
                    return makeOSProcess(cmd.Process), nil
                },
            },
        },
    }
}

func makeOSProcess(proc *os.Process) *tengopkg.ImmutableMap {
    return &tengopkg.ImmutableMap{
        Value: map[string]tengopkg.Object{
            "kill": &tengopkg.UserFunction{
                Name:  "kill",
                Value: stdlib.FuncARE(proc.Kill),
            },
            "release": &tengopkg.UserFunction{
                Name:  "release",
                Value: stdlib.FuncARE(proc.Release),
            },
            "signal": &tengopkg.UserFunction{
                Name: "signal",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 1 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }
                    i1, ok := tengopkg.ToInt64(args[0])
                    if !ok {
                        return nil, tengopkg.ErrInvalidArgumentType{
                            Name:     "first",
                            Expected: "int(compatible)",
                            Found:    args[0].TypeName(),
                        }
                    }
                    return tengolib.WrapError(proc.Signal(syscall.Signal(i1))), nil
                },
            },
            "wait": &tengopkg.UserFunction{
                Name: "wait",
                Value: func(args ...tengopkg.Object) (tengopkg.Object, error) {
                    if len(args) != 0 {
                        return nil, tengopkg.ErrWrongNumArguments
                    }
                    state, err := proc.Wait()
                    if err != nil {
                        return tengolib.WrapError(err), nil
                    }
                    return makeOSProcessState(state), nil
                },
            },
        },
    }
}

func makeOSProcessState(state *os.ProcessState) *tengopkg.ImmutableMap {
    return &tengopkg.ImmutableMap{
        Value: map[string]tengopkg.Object{
            "exited": &tengopkg.UserFunction{
                Name:  "exited",
                Value: stdlib.FuncARB(state.Exited),
            },
            "pid": &tengopkg.UserFunction{
                Name:  "pid",
                Value: stdlib.FuncARI(state.Pid),
            },
            "string": &tengopkg.UserFunction{
                Name:  "string",
                Value: stdlib.FuncARS(state.String),
            },
            "success": &tengopkg.UserFunction{
                Name:  "success",
                Value: stdlib.FuncARB(state.Success),
            },
        },
    }
}