modules-gopwsh/shared.go

package gopwsh

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strings"
)

const newline = "\r\n"

type Shell interface {
	Execute(cmd string) (string, error)
	Exit()
}

type Waiter interface {
	Wait() error
}

type shell struct {
	handle Waiter
	stdin  io.Writer
	stdout io.Reader
}

func (e PowershellError) Error() string {
	return e.Message
}

type PowershellError struct {
	Status   int    `json:"Status"`
	Response any    `json:"Response"`
	Message  string `json:"Message"`
	Data     struct {
	} `json:"Data"`
	InnerException any    `json:"InnerException"`
	StackTrace     string `json:"StackTrace"`
	HelpLink       any    `json:"HelpLink"`
	Source         string `json:"Source"`
	HResult        int    `json:"HResult"`
}

var noShellError = errors.New(
	"Cannot execute commands on closed shells.",
)

func newPowershellError(stdErr string) error {

	psErr := &PowershellError{}
	_ = json.Unmarshal([]byte(stdErr), psErr)
	return psErr
}

func (s *shell) Execute(cmd string) (string, error) {
	if s.handle == nil {
		return "", noShellError
	}

	end := createBoundary()

	command := fmt.Sprintf(
		"try{%s}catch{$psitem.Exception | convertto-json -WarningAction SilentlyContinue;Write-Output '$$ERROR$$'}",
		cmd,
	)

	full := fmt.Sprintf(
		"%s; Write-Output '%s'%s",
		command,
		end,
		newline,
	)

	_, err := s.stdin.Write([]byte(full))
	if err != nil {
		return "", err
	}

	stdOut := ""

	err = reader(s.stdout, end, &stdOut)
	if err != nil {
		return "", err
	}

	return stdOut, nil
}

func (s *shell) Exit() {
	_, _ = s.stdin.Write([]byte("exit" + newline))

	closer, ok := s.stdin.(io.Closer)
	if ok {
		closer.Close()
	}

	s.handle.Wait()

	s.handle = nil
	s.stdin = nil
	s.stdout = nil
}

func reader(
	stream io.Reader,
	boundary string,
	buffer *string,
) error {
	output := ""
	bufsize := 64
	marker := boundary + newline

	for {
		buf := make([]byte, bufsize)
		read, err := stream.Read(buf)
		if err != nil {
			return err
		}

		if strings.HasSuffix(string(buf[:read]), "$$ERROR$$") {
			err := newPowershellError(output)
			return err
		}

		output = output + string(buf[:read])

		if strings.HasSuffix(output, marker) {
			break
		}
	}

	*buffer = strings.TrimSuffix(output, marker)

	return nil
}

func createBoundary() string {
	c := 12
	b := make([]byte, c)

	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}

	randomString := hex.EncodeToString(b)

	return "$$" + randomString + "$$"
}

func StartProcess(
	cmd string,
	args ...string,
) (
	Waiter,
	io.Writer,
	io.Reader,
	error,
) {
	command := exec.Command(cmd, args...)

	stdin, err := command.StdinPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	stdout, err := command.StdoutPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	err = command.Start()
	if err != nil {
		return nil, nil, nil, err
	}

	return command, stdin, stdout, nil
}