1package policy
  2
  3import (
  4	"bytes"
  5	"context"
  6	"errors"
  7	"fmt"
  8	"os"
  9	"os/exec"
 10	"time"
 11
 12	"go.starlark.net/starlark"
 13)
 14
 15// ExecModule provides the ability to execute external binary commands from Starlark policies.
 16type ExecModule struct{}
 17
 18func (m *ExecModule) String() string       { return "<module exec>" }
 19func (m *ExecModule) Type() string         { return "module" }
 20func (m *ExecModule) Freeze()              {}
 21func (m *ExecModule) Truth() starlark.Bool { return true }
 22func (m *ExecModule) Hash() (uint32, error) {
 23	return 0, errors.New("unhashable")
 24}
 25
 26func (m *ExecModule) Attr(name string) (starlark.Value, error) {
 27	switch name {
 28	case "pipe_to_process":
 29		return starlark.NewBuiltin("pipe_to_process", m.pipeToProcess), nil
 30	}
 31	return nil, nil
 32}
 33
 34func (m *ExecModule) AttrNames() []string {
 35	return []string{"pipe_to_process"}
 36}
 37
 38// pipeToProcess executes an external binary, piping input_data to stdin,
 39// and returns an ExecResult with code, stdout, stderr, and timed_out fields.
 40func (m *ExecModule) pipeToProcess(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 41	var command string
 42	var inputData starlark.Value = starlark.None
 43	var argList starlark.Value = starlark.None
 44	var timeoutVal starlark.Value = starlark.None
 45	var envDict starlark.Value = starlark.None
 46
 47	if err := starlark.UnpackArgs(b.Name(), args, kwargs,
 48		"command", &command,
 49		"input_data?", &inputData,
 50		"args?", &argList,
 51		"timeout?", &timeoutVal,
 52		"environment?", &envDict,
 53	); err != nil {
 54		return nil, err
 55	}
 56
 57	var inputBytes []byte
 58	switch v := inputData.(type) {
 59	case starlark.String:
 60		inputBytes = []byte(v)
 61	case starlark.Bytes:
 62		inputBytes = []byte(v)
 63	case starlark.NoneType:
 64		inputBytes = nil
 65	default:
 66		return nil, fmt.Errorf("input_data must be a string or bytes, got %s", inputData.Type())
 67	}
 68
 69	var goArgs []string
 70	if argList != starlark.None {
 71		lst, ok := argList.(*starlark.List)
 72		if !ok {
 73			return nil, fmt.Errorf("args must be a list, got %s", argList.Type())
 74		}
 75		for it := range lst.Elements() {
 76			s, ok := it.(starlark.String)
 77			if !ok {
 78				return nil, fmt.Errorf("args must be a list of strings, got %s", it.Type())
 79			}
 80			goArgs = append(goArgs, string(s))
 81		}
 82	}
 83
 84	var timeoutDur = 30 * time.Second
 85	if timeoutVal != starlark.None {
 86		switch v := timeoutVal.(type) {
 87		case starlark.Int:
 88			if i, ok := v.Int64(); ok {
 89				timeoutDur = time.Duration(i) * time.Second
 90			} else {
 91				return nil, errors.New("timeout must be a valid integer")
 92			}
 93		case starlark.Float:
 94			f, _ := starlark.AsFloat(v)
 95			timeoutDur = time.Duration(f) * time.Second
 96		default:
 97			return nil, fmt.Errorf("timeout must be a number, got %s", timeoutVal.Type())
 98		}
 99	}
100
101	hasPath := false
102	var env []string
103	if envDict != starlark.None {
104		dict, ok := envDict.(*starlark.Dict)
105		if !ok {
106			return nil, fmt.Errorf("environment must be a dict, got %s", envDict.Type())
107		}
108		for _, item := range dict.Items() {
109			key, ok := item[0].(starlark.String)
110			if !ok {
111				return nil, fmt.Errorf("environment key must be a string, got %s", item[0].Type())
112			}
113			val, ok := item[1].(starlark.String)
114			if !ok {
115				return nil, fmt.Errorf("environment value must be a string, got %s", item[1].Type())
116			}
117			env = append(env, string(key)+"="+string(val))
118			if key == "PATH" && len(val) > 0 {
119				hasPath = true
120			}
121		}
122	}
123
124	if !hasPath {
125		if path := os.Getenv("PATH"); path != "" {
126			env = append(env, "PATH="+path)
127		}
128	}
129
130	ctx, ok := thread.Local("context").(context.Context)
131	if !ok || ctx == nil {
132		ctx = context.Background()
133	}
134	ctx, cancel := context.WithTimeout(ctx, timeoutDur)
135	defer cancel()
136
137	cmd := exec.CommandContext(ctx, command, goArgs...)
138	cmd.Env = env
139	cmd.Stdin = bytes.NewReader(inputBytes)
140
141	var stdout, stderr bytes.Buffer
142	cmd.Stdout = &stdout
143	cmd.Stderr = &stderr
144
145	runErr := cmd.Run()
146
147	timedOut := errors.Is(ctx.Err(), context.DeadlineExceeded)
148
149	code := 0
150	if runErr != nil {
151		if timedOut {
152			code = -1
153		} else if exitErr, ok := errors.AsType[*exec.ExitError](runErr); ok {
154			code = exitErr.ExitCode()
155		} else {
156			return nil, fmt.Errorf("exec failed: %w", runErr)
157		}
158	}
159
160	return &ExecResult{
161		code:     code,
162		stdout:   stdout.String(),
163		stderr:   stderr.String(),
164		timedOut: timedOut,
165	}, nil
166}
167
168// ExecResult represents the result of an external process execution.
169type ExecResult struct {
170	code     int
171	stdout   string
172	stderr   string
173	timedOut bool
174}
175
176func (r *ExecResult) String() string {
177	return fmt.Sprintf("<ExecResult code=%d>", r.code)
178}
179func (r *ExecResult) Type() string { return "ExecResult" }
180func (r *ExecResult) Freeze()      {}
181func (r *ExecResult) Truth() starlark.Bool {
182	return starlark.Bool(r.code == 0 && !r.timedOut)
183}
184func (r *ExecResult) Hash() (uint32, error) {
185	return 0, errors.New("unhashable")
186}
187
188func (r *ExecResult) Attr(name string) (starlark.Value, error) {
189	switch name {
190	case "code":
191		return starlark.MakeInt(r.code), nil
192	case "stdout":
193		return starlark.String(r.stdout), nil
194	case "stderr":
195		return starlark.String(r.stderr), nil
196	case "timed_out":
197		return starlark.Bool(r.timedOut), nil
198	}
199	return nil, nil
200}
201
202func (r *ExecResult) AttrNames() []string {
203	return []string{"code", "stdout", "stderr", "timed_out"}
204}