創建可交互進程

go language 版本:

需要依賴winpty.dll 和 winpty-agent.exe

winpty.go

package winpty

import (
	"fmt"
	"io"
	"os"
	"syscall"
	"unsafe"
)

type Options struct {
	// DllDir is the path to winpty.dll and winpty-agent.exe
	DllDir string
	// FilePath sets the title of the console
	FilePath string
	// Command is the full command to launch
	Command string
	// Dir sets the current working directory for the command
	Dir string
	// Env sets the environment variables. Use the format VAR=VAL.
	Env []string
	// AgentFlags to pass to agent config creation
	AgentFlags uint64
	SpawnFlag uint32
	MouseModes int
	// Initial size for Columns and Rows
	InitialCols uint32
	InitialRows uint32
	agentTimeoutMs *uint64
}

type WinPTY struct {
	Stdin         *os.File
	Stdout        *os.File
	Stderr        *os.File
	pty           uintptr
	winptyConfigT uintptr
	procHandle    uintptr
	closed        bool
	exitCode      *int
}

func (pty *WinPTY) Pid() int {
	pid, _, _ :=GetProcessId.Call(pty.procHandle)
	return int(pid)
}
//這裏不能講一個file結構體賦值給一個WriteCloser的原因是file結構體的Close方法的第一個參數是一個file指針而不是file指針,也就是說接口方法的對應的結構體方法的第一個參數可以是結構體對或者結構體指針
func (pty *WinPTY) GetStdin() io.Reader {
	return pty.Stdin//這裏的類型是機構體指針還是結構體本身取決於該結構體實現的接口方法的第一個參數是指針還是結構體
}

func (pty *WinPTY) GetStdout() io.Writer {
	return pty.Stdout
}

func (pty *WinPTY) GetStderr() io.Writer {
	return pty.Stderr
}

// the same as open, but uses defaults for Env
func OpenPTY(dllPrefix, cmd, dir string,isColor bool) (*WinPTY, error) {
	var flag uint64 =WINPTY_FLAG_PLAIN_OUTPUT
	if isColor{
		flag = WINPTY_FLAG_COLOR_ESCAPES
	}
	flag = flag|WINPTY_FLAG_ALLOW_CURPROC_DESKTOP_CREATION
	return CreateProcessWithOptions(Options{
		DllDir:     dllPrefix,
		Command:    cmd,
		Dir:        dir,
		Env:        os.Environ(),
		AgentFlags: flag,
	})
}

func setOptsDefaultValues(options *Options){
	// Set the initial size to 40x40 if options is 0
	if options.InitialCols <= 0 {
		options.InitialCols = 40
	}
	if options.InitialRows <= 0 {
		options.InitialRows = 40
	}
	if options.agentTimeoutMs == nil {
		t:=uint64(syscall.INFINITE)
		options.agentTimeoutMs = &t
	}
	if options.SpawnFlag == 0{
		options.SpawnFlag = 1
	}
	if options.MouseModes<0{
		options.MouseModes = 0
	}
}

func CreateProcessWithOptions(options Options) (*WinPTY, error) {
	setOptsDefaultValues(&options)
	setupDefines(options.DllDir)
	// create config with specified AgentFlags
	winptyConfigT, err := createAgentCfg(options.AgentFlags)
	if err != nil {
		return nil, err
	}

	winpty_config_set_initial_size.Call(winptyConfigT, uintptr(options.InitialCols), uintptr(options.InitialRows))
	SetMouseMode(winptyConfigT,options.MouseModes)

	var openErr uintptr
	defer winpty_error_free.Call(openErr)
	pty, _, _ := winpty_open.Call(winptyConfigT, uintptr(unsafe.Pointer(openErr)))

	if pty == uintptr(0) {
		return nil, fmt.Errorf("Error Launching WinPTY agent, %s ", GetErrorMessage(openErr))
	}

	SetAgentTimeout(winptyConfigT,*options.agentTimeoutMs)
	winpty_config_free.Call(winptyConfigT)

	stdinName, _, _ := winpty_conin_name.Call(pty)
	stdoutName, _, _ := winpty_conout_name.Call(pty)
	stderrName, _, _ := winpty_conerr_name.Call(pty)

	obj := &WinPTY{}

	stdinHandle, err := syscall.CreateFile((*uint16)(unsafe.Pointer(stdinName)), syscall.GENERIC_WRITE, 0, nil, syscall.OPEN_EXISTING, 0, 0)
	if err != nil {
		return nil, fmt.Errorf("Error getting stdin handle. %s ", err)
	}
	obj.Stdin = os.NewFile(uintptr(stdinHandle), "stdin")

	stdoutHandle, err := syscall.CreateFile((*uint16)(unsafe.Pointer(stdoutName)), syscall.GENERIC_READ, 0, nil, syscall.OPEN_EXISTING, 0, 0)
	if err != nil {
		return nil, fmt.Errorf("Error getting stdout handle. %s ", err)
	}
	obj.Stdout = os.NewFile(uintptr(stdoutHandle), "stdout")

	if options.AgentFlags&WINPTY_FLAG_CONERR == WINPTY_FLAG_CONERR{
		stderrHandle, err := syscall.CreateFile((*uint16)(unsafe.Pointer(stderrName)), syscall.GENERIC_READ, 0, nil, syscall.OPEN_EXISTING, 0, 0)
		if err != nil {
			return nil, fmt.Errorf("Error getting stderr handle. %s ", err)
		}
		obj.Stderr = os.NewFile(uintptr(stderrHandle), "stderr")
	}

	spawnCfg, err := createSpawnCfg(options.SpawnFlag, options.FilePath, options.Command, options.Dir, options.Env)
	if err != nil {
		return nil, err
	}
	var (
		spawnErr  uintptr
		lastError *uint32
	)
	spawnRet, _, _ := winpty_spawn.Call(pty, spawnCfg, uintptr(unsafe.Pointer(&obj.procHandle)), uintptr(0), uintptr(unsafe.Pointer(lastError)), uintptr(unsafe.Pointer(spawnErr)))
	_, _, _ = winpty_spawn_config_free.Call(spawnCfg)
	defer winpty_error_free.Call(spawnErr)

	if spawnRet == 0 {
		return nil, fmt.Errorf("Error spawning process... ")
	} else {
		obj.pty = pty
		return obj, nil
	}
}

//設置窗口大小
func (pty *WinPTY) SetSize(wsCol, wsRow uint32) {
	if wsCol == 0 || wsRow == 0 {
		return
	}
	winpty_set_size.Call(pty.pty, uintptr(wsCol), uintptr(wsRow), uintptr(0))
}
//關閉進程
func (pty *WinPTY) Close(){
	if pty.closed {
		return
	}
	winpty_free.Call(pty.pty)
	_ = pty.Stdin.Close()
	_ = pty.Stdout.Close()
	_ = syscall.CloseHandle(syscall.Handle(pty.procHandle))
	pty.closed = true
}

//等待進程結束
func (pty *WinPTY)Wait() error  {
	err:=WaitForSingleObject(pty.procHandle,syscall.INFINITE)
	pty.Close()
	return err
}
//獲取創建的目標進程的句柄
func (pty *WinPTY) GetProcHandle() uintptr {
	return pty.procHandle
}
//獲取代理進程的句柄
func (pty *WinPTY)GetAgentProcHandle() uintptr  {
	agentProcH,_,_:=winpty_agent_process.Call(pty.pty)
	return agentProcH
}
//設置代理程序啓動的等待時間
func SetAgentTimeout(winptyConfigT uintptr,timeoutMs uint64) {
	winpty_config_set_agent_timeout.Call(winptyConfigT, uintptr(timeoutMs))
}
//設置鼠標的模式
func SetMouseMode(winptyConfigT uintptr,mode int)  {
	winpty_config_set_mouse_mode.Call(winptyConfigT, uintptr(mode))
}
//獲取退出代碼
func (pty *WinPTY)ExitCode() int  {
	if pty.exitCode == nil{
		code,err:=GetExitCodeProcess(pty.procHandle)
		if err!=nil{
			code = -1
		}
		pty.exitCode = &code
	}
	return *pty.exitCode
}
//強制殺掉進程
func (pty *WinPTY)Kill()  {
	code:=-1
	_ = syscall.TerminateProcess(syscall.Handle(pty.procHandle), uint32(code))
	pty.Close()
}
winpty_amd64.go
package winpty

import (
	"fmt"
	"syscall"
	"unsafe"
)

func createAgentCfg(flags uint64) (uintptr, error) {
	var errorPtr uintptr
	defer winpty_error_free.Call(errorPtr)

	winptyConfigT, _, _ := winpty_config_new.Call(uintptr(flags), uintptr(unsafe.Pointer(errorPtr)))
	if winptyConfigT == uintptr(0) {
		return 0, fmt.Errorf("Unable to create agent config, %s ", GetErrorMessage(errorPtr))
	}

	return winptyConfigT, nil
}

func createSpawnCfg(flags uint32, filePath, cmdline, cwd string, env []string) (uintptr, error) {
	var errorPtr uintptr
	defer winpty_error_free.Call(errorPtr)

	cmdLineStr, err := syscall.UTF16PtrFromString(cmdline)
	if err != nil {
		return 0, fmt.Errorf("Failed to convert cmd to pointer. ")
	}

	filepath, err := syscall.UTF16PtrFromString(filePath)
	if err != nil {
		return 0, fmt.Errorf("Failed to convert app name to pointer. ")
	}

	cwdStr, err := syscall.UTF16PtrFromString(cwd)
	if err != nil {
		return 0, fmt.Errorf("Failed to convert working directory to pointer. ")
	}

	envStr, err := UTF16PtrFromStringArray(env)

	if err != nil {
		return 0, fmt.Errorf("Failed to convert cmd to pointer. ")
	}

	spawnCfg, _, _ := winpty_spawn_config_new.Call(
		uintptr(flags),
		uintptr(unsafe.Pointer(filepath)),
		uintptr(unsafe.Pointer(cmdLineStr)),
		uintptr(unsafe.Pointer(cwdStr)),
		uintptr(unsafe.Pointer(envStr)),
		uintptr(unsafe.Pointer(errorPtr)),
	)

	if spawnCfg == uintptr(0) {
		return 0, fmt.Errorf("Unable to create spawn config, %s ", GetErrorMessage(errorPtr))
	}

	return spawnCfg, nil
}

util.go

package winpty

import (
	"errors"
	"syscall"
	"unicode/utf16"
	"unsafe"
)

func UTF16PtrToString(p *uint16) string {
	var (
		sizeTest uint16
		finalStr = make([]uint16, 0)
	)
	for {
		if *p == uint16(0) {
			break
		}

		finalStr = append(finalStr, *p)
		p = (*uint16)(unsafe.Pointer(uintptr(unsafe.Pointer(p)) + unsafe.Sizeof(sizeTest)))
	}
	return string(utf16.Decode(finalStr[0:]))
}

func UTF16PtrFromStringArray(s []string) (*uint16, error) {
	var r []uint16

	for _, ss := range s {
		a, err := syscall.UTF16FromString(ss)
		if err != nil {
			return nil, err
		}

		r = append(r, a...)
	}

	r = append(r, 0)

	return &r[0], nil
}

func GetErrorMessage(err uintptr) string {
	msgPtr, _, _ := winpty_error_msg.Call(err)
	if msgPtr == uintptr(0) {
		return "Unknown Error"
	}
	return UTF16PtrToString((*uint16)(unsafe.Pointer(msgPtr)))
}

func GetErrorCode(err uintptr) uint32  {
	code,_,_:=winpty_error_code.Call(err)
	return uint32(code)
}


func GetExitCodeProcess(procHandle uintptr) (int,error){
	getExitCode :=kernel32.NewProc("GetExitCodeProcess")
	var code uintptr
	ok,_,_:=getExitCode.Call(procHandle,uintptr(unsafe.Pointer(&code)))
	if ok == 0{
		err:=syscall.GetLastError()
		return -1,err
	}
	return int(code),nil
}


var TimeoutError = errors.New("timeout error")
var InvalidHandleError = errors.New("invalid handle")

func WaitForSingleObject(procHandle uintptr,milliseconds int64) error {
	wait:=kernel32.NewProc("WaitForSingleObject")
	res,_,_:=wait.Call(procHandle, uintptr(milliseconds))
	if res == 0{
		return nil
	}else if res == 258{
		return TimeoutError
	}else {
		return InvalidHandleError
	}
}


defines.go

package winpty

import (
	"path/filepath"
	"syscall"
)

const (
	WINPTY_SPAWN_FLAG_AUTO_SHUTDOWN       = 1
	WINPTY_SPAWN_FLAG_EXIT_AFTER_SHUTDOWN = 2

	WINPTY_FLAG_CONERR = 0x1
	WINPTY_FLAG_PLAIN_OUTPUT = 0x2
	WINPTY_FLAG_COLOR_ESCAPES = 0x4
	WINPTY_FLAG_ALLOW_CURPROC_DESKTOP_CREATION = 0x8

	WINPTY_MOUSE_MODE_NONE  = 0
	WINPTY_MOUSE_MODE_AUTO  = 1
	WINPTY_MOUSE_MODE_FORCE = 2
)

var (
	modWinPTY *syscall.LazyDLL
	kernel32 *syscall.LazyDLL
	// Error handling...
	winpty_error_code *syscall.LazyProc
	winpty_error_msg  *syscall.LazyProc
	winpty_error_free *syscall.LazyProc

	// Configuration of a new agent.
	winpty_config_new               *syscall.LazyProc
	winpty_config_free              *syscall.LazyProc
	winpty_config_set_initial_size  *syscall.LazyProc
	winpty_config_set_mouse_mode    *syscall.LazyProc
	winpty_config_set_agent_timeout *syscall.LazyProc

	// Start the agent.
	winpty_open          *syscall.LazyProc
	winpty_agent_process *syscall.LazyProc

	// I/O Pipes
	winpty_conin_name  *syscall.LazyProc
	winpty_conout_name *syscall.LazyProc
	winpty_conerr_name *syscall.LazyProc

	// Agent RPC Calls
	winpty_spawn_config_new  *syscall.LazyProc
	winpty_spawn_config_free *syscall.LazyProc
	winpty_spawn             *syscall.LazyProc
	winpty_set_size          *syscall.LazyProc
	winpty_free              *syscall.LazyProc

	//windows api
	GetProcessId			*syscall.LazyProc
)

func init() {
	kernel32 = syscall.NewLazyDLL("kernel32.dll")
	GetProcessId = kernel32.NewProc("GetProcessId")
}

func setupDefines(dllPrefix string) {

	if modWinPTY != nil {
		return
	}

	modWinPTY = syscall.NewLazyDLL(filepath.Join(dllPrefix, `winpty.dll`))
	// Error handling...
	winpty_error_code = modWinPTY.NewProc("winpty_error_code")
	winpty_error_msg = modWinPTY.NewProc("winpty_error_msg")
	winpty_error_free = modWinPTY.NewProc("winpty_error_free")

	// Configuration of a new agent.
	winpty_config_new = modWinPTY.NewProc("winpty_config_new")
	winpty_config_free = modWinPTY.NewProc("winpty_config_free")
	winpty_config_set_initial_size = modWinPTY.NewProc("winpty_config_set_initial_size")
	winpty_config_set_mouse_mode = modWinPTY.NewProc("winpty_config_set_mouse_mode")
	winpty_config_set_agent_timeout = modWinPTY.NewProc("winpty_config_set_agent_timeout")

	// Start the agent.
	winpty_open = modWinPTY.NewProc("winpty_open")
	winpty_agent_process = modWinPTY.NewProc("winpty_agent_process")

	// I/O Pipes
	winpty_conin_name = modWinPTY.NewProc("winpty_conin_name")
	winpty_conout_name = modWinPTY.NewProc("winpty_conout_name")
	winpty_conerr_name = modWinPTY.NewProc("winpty_conerr_name")

	// Agent RPC Calls
	winpty_spawn_config_new = modWinPTY.NewProc("winpty_spawn_config_new")
	winpty_spawn_config_free = modWinPTY.NewProc("winpty_spawn_config_free")
	winpty_spawn = modWinPTY.NewProc("winpty_spawn")
	winpty_set_size = modWinPTY.NewProc("winpty_set_size")
	winpty_free = modWinPTY.NewProc("winpty_free")
}

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章