创建可交互进程

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

 

你可能感兴趣的:(golang)