initial commit

shared.go

@@ -0,0 +1,182 @@
+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
+}