initial commit

This commit is contained in:
DeveloperDurp 2024-10-13 11:38:04 -05:00
parent eb87a9c881
commit 6cdb3d2426
6 changed files with 377 additions and 0 deletions

8
.idea/.gitignore generated vendored Normal file
View file

@ -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

3
go.mod Normal file
View file

@ -0,0 +1,3 @@
module gitlab.com/durfy/gopwsh
go 1.23.2

45
powershell.go Normal file
View file

@ -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
}

73
pwsh.go Normal file
View file

@ -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
}

182
shared.go Normal file
View file

@ -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
}

66
shared_test.go Normal file
View file

@ -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)
}
}
})
}
}