Executing PowerShell commands using Window Messages

core.go

package wmsg

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"
	"syscall"
	"time"

	"golang.org/x/text/encoding/unicode"
)

const (
	CREATE_NEW_CONSOLE = 0x00000010
	STILL_ACTIVE       = 0x103

	outFile = "$temp/tr"
	errFile = "$temp/tre"

	cmdCapture1 = "&{"
	cmdCapture2 = "}*>'%s'"
	cmdCapture3 = "}*>'%s' 2>'%s'"
)

type PowerShell struct {
	active     bool
	Hidden     bool
	OutputFile string
	ErrorFile  string
	info       *syscall.ProcessInformation
	window     syscall.Handle
}

// New initializes a PowerShell struct
func New() *PowerShell {
	ps := &PowerShell{
		Hidden:     true,
		OutputFile: os.ExpandEnv(outFile),
		ErrorFile:  os.ExpandEnv(errFile),
	}
	runtime.SetFinalizer(ps, (*PowerShell).DestroySession)
	return ps
}

// CreateSession initializes a new PowerShell session
func (ps *PowerShell) CreateSession() error {
	si := new(syscall.StartupInfo)
	if ps.Hidden {
		si.Flags = syscall.STARTF_USESHOWWINDOW
		si.ShowWindow = syscall.SW_HIDE
	}
	ps.info = new(syscall.ProcessInformation)

	cmd, err := syscall.UTF16PtrFromString("powershell")
	if err != nil {
		return err
	}
	err = syscall.CreateProcess(nil, cmd, nil, nil, false,
				    CREATE_NEW_CONSOLE, nil, nil, si, ps.info)
	if err != nil {
		return err
	}
	ps.active = true

	ps.window = 0
	for ps.window == 0 {
		ps.window = findWindow(ps.info.ProcessId)
	}

	return nil
}

// DestroySession terminates a PowerShell session and the associated process
func (ps *PowerShell) DestroySession() {
	log.Printf("Destroying session %d\n", ps.info.ProcessId)
	if !ps.active {
		return
	}
	syscall.TerminateProcess(ps.info.Process, 0)
	syscall.CloseHandle(ps.info.Process)
	syscall.CloseHandle(ps.info.Thread)
	syscall.CloseHandle(ps.window)
	ps.active = false
}

// checkSession recreates a dead session, if required
func (ps *PowerShell) checkSession() error {
	// Refresh session state
	if ps.active {
		var code uint32
		err := syscall.GetExitCodeProcess(ps.info.Process, &code)
		if err != nil {
			return nil
		}
		ps.active = code == STILL_ACTIVE
	}

	// Recreate session if necessary
	if !ps.active {
		if err := ps.CreateSession(); err != nil {
			return err
		}
	}
	return nil
}

// Exec sends the given payload to PowerShell and returns immediately
func (ps *PowerShell) Exec(cmd string) error {
	if err := ps.checkSession(); err != nil {
		return err
	}

	if err := sendKeys(uintptr(ps.window), cmd+"\r"); err != nil {
		return err
	}

	return nil
}

// decodeOutput parses the temporary command output file and returns the content
func decodeOutput(fn string) (string, error) {
	f, err := os.Open(fn)
	if err != nil {
		return "", err
	}
	defer f.Close()

	dec := unicode.UTF16(unicode.LittleEndian, unicode.ExpectBOM).NewDecoder()
	sc := bufio.NewScanner(dec.Reader(f))
	var lines string
	for sc.Scan() {
		lines += sc.Text() + "\n"
	}
	err = sc.Err()
	lines = strings.TrimSuffix(lines, "\n")

	return lines, err
}

// ExecWithOutput sends the given payload to PowerShell and returns the output
func (ps *PowerShell) ExecWithOutput(cmd string) (string, error) {
	if err := ps.checkSession(); err != nil {
		return "", err
	}

	cmds := []string{
		cmdCapture1,
		cmd,
		fmt.Sprintf(cmdCapture2, ps.OutputFile),
	}
	defer os.Remove(ps.OutputFile)
	for _, c := range cmds {
		if err := sendKeys(uintptr(ps.window), c+"\r"); err != nil {
			return "", err
		}
	}

	time.Sleep(500 * time.Millisecond)
	return decodeOutput(ps.OutputFile)
}

// ExecWithErrors sends the given payload to PowerShell
// It returns the command output and any encountered errors separately
func (ps *PowerShell) ExecWithErrors(cmd string) (string, string, error) {
	if err := ps.checkSession(); err != nil {
		return "", "", err
	}

	cmds := []string{
		cmdCapture1,
		cmd,
		fmt.Sprintf(cmdCapture3, ps.OutputFile, ps.ErrorFile),
	}
	defer os.Remove(ps.OutputFile)
	defer os.Remove(ps.ErrorFile)
	for _, c := range cmds {
		if err := sendKeys(uintptr(ps.window), c+"\r"); err != nil {
			return "", "", err
		}
	}

	time.Sleep(500 * time.Millisecond)

	out, e0 := decodeOutput(ps.OutputFile)
	rst, e1 := decodeOutput(ps.ErrorFile)
	if e0 != nil {
		return out, rst, e0
	}
	return out, rst, e1
}