From 6cdb3d24266dde6907ac5ea373d5614b9dee28ca Mon Sep 17 00:00:00 2001 From: DeveloperDurp Date: Sun, 13 Oct 2024 11:38:04 -0500 Subject: [PATCH] initial commit --- .idea/.gitignore | 8 +++ go.mod | 3 + powershell.go | 45 ++++++++++++ pwsh.go | 73 +++++++++++++++++++ shared.go | 182 +++++++++++++++++++++++++++++++++++++++++++++++ shared_test.go | 66 +++++++++++++++++ 6 files changed, 377 insertions(+) create mode 100644 .idea/.gitignore create mode 100644 go.mod create mode 100644 powershell.go create mode 100644 pwsh.go create mode 100644 shared.go create mode 100644 shared_test.go diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..13566b8 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,8 @@ +# Default ignored files +/shelf/ +/workspace.xml +# Editor-based HTTP Client requests +/httpRequests/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..1d82735 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module gitlab.com/durfy/gopwsh + +go 1.23.2 diff --git a/powershell.go b/powershell.go new file mode 100644 index 0000000..d6a9f34 --- /dev/null +++ b/powershell.go @@ -0,0 +1,45 @@ +package gopwsh + +type Powershell struct{} + +func NewPowershellHandler() Powershell { + return Powershell{} +} + +func RunPowershellCommand(cmd string) (string, error) { + + ps := NewPowershellHandler() + shell, err := ps.New() + if err != nil { + return "", err + } + + defer shell.Exit() + + sout, err := shell.Execute(cmd) + if err != nil { + return sout, err + } + + return sout, nil +} + +func (ps *Powershell) New() (Shell, error) { + + handle, stdIn, stdOut, err := StartProcess( + "powershell.exe", + "-NoExit", + "-NoProfile", + "-Command", + "-", + ) + if err != nil { + return nil, err + } + + return &shell{ + handle, + stdIn, + stdOut, + }, nil +} diff --git a/pwsh.go b/pwsh.go new file mode 100644 index 0000000..445869e --- /dev/null +++ b/pwsh.go @@ -0,0 +1,73 @@ +package gopwsh + +import ( + "os/exec" + "runtime" +) + +type Pwsh struct{} + +func NewPwshHandler() Pwsh { + return Pwsh{} +} + +func getPwshBinary() string { + switch runtime.GOOS { + case "windows": + return "pwsh.exe" + default: + return "pwsh" + } +} + +func IsPwshInstalled() bool { + binary := getPwshBinary() + cmd := exec.Command(binary) + output, err := cmd.CombinedOutput() + if err != nil { + return false + } + if string(output) == "" { + return false + } + return true +} + +func RunPwshCommand(cmd string) (string, error) { + + ps := NewPwshHandler() + shell, err := ps.New() + if err != nil { + return "", err + } + + defer shell.Exit() + + stdOut, err := shell.Execute(cmd) + if err != nil { + return stdOut, err + } + + return stdOut, nil +} + +func (ps *Pwsh) New() (Shell, error) { + + binary := getPwshBinary() + handle, stdIn, stdOut, err := StartProcess( + binary, + "-NoExit", + "-NoProfile", + "-Command", + "-", + ) + if err != nil { + return nil, err + } + + return &shell{ + handle, + stdIn, + stdOut, + }, nil +} diff --git a/shared.go b/shared.go new file mode 100644 index 0000000..ffaf3bd --- /dev/null +++ b/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 +} diff --git a/shared_test.go b/shared_test.go new file mode 100644 index 0000000..1f27056 --- /dev/null +++ b/shared_test.go @@ -0,0 +1,66 @@ +package gopwsh + +import ( + "errors" + "testing" +) + +func TestRunPowershellCommand(t *testing.T) { + tests := []struct { + name string + cmd string + expectedErr error + }{ + { + name: "test error", + cmd: "irm tgssdcfas.asdfasdf", + expectedErr: errors.New("The remote name could not be resolved: 'tgssdcfas.asdfasdf'"), + }, + { + name: "Test working command", + cmd: `get-host`, + expectedErr: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := RunPowershellCommand(tt.cmd) + if err != nil { + if err.Error() != tt.expectedErr.Error() { + t.Errorf("RunPowershellCommand() error = %v, want error %v", err, tt.expectedErr) + } + } + }) + } +} + +func TestRunPwshCommand(t *testing.T) { + tests := []struct { + name string + cmd string + expectedErr error + }{ + { + name: "test error", + cmd: "irm tgssdcfas.asdfasdf", + expectedErr: errors.New("No such host is known. (tgssdcfas.asdfasdf:80)"), + }, + { + name: "Test working command", + cmd: `get-host`, + expectedErr: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := RunPwshCommand(tt.cmd) + if err != nil { + if err.Error() != tt.expectedErr.Error() { + t.Errorf("RunPwshCommand() error = %v, want error %v", err, tt.expectedErr) + } + } + }) + } +}