gitee.com/h79/goutils@v1.22.10/common/system/cmd.go

package system

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

type Config struct {
	Stdin     io.Reader
	Stdout    io.Writer
	Stderr    io.Writer
	EnableStd bool `json:"enableStd" yaml:"enableStd" xml:"enableStd"`
	EnableEnv bool `json:"enableEnv" yaml:"enableEnv" xml:"enableEnv"`
}

type Result struct {
	stdout    bytes.Buffer
	stderr    bytes.Buffer
	Command   string
	Err       error
	StartTime time.Time
	EndTime   time.Time
}

func (r Result) Output() string {
	if r.Err != nil {
		return r.stderr.String()
	}
	return r.stdout.String()
}

func (r Result) Error() string {
	if r.Err != nil {
		return fmt.Sprintf("Result: %s, err: %v", r.Command, r.Err)
	}
	return fmt.Sprintf("Result: %s, Result: %s, StartTime:%v, EndTime: %v", r.Command, r.Output(), r.StartTime, r.EndTime)
}

func (c *Cmd) init(conf *Config, env ...string) *exec.Cmd {
	if conf == nil {
		return c.Cmd
	}
	if conf.EnableEnv {
		c.Cmd.Env = os.Environ()
		c.Cmd.Env = append(c.Cmd.Env, env...)
	}
	if conf.EnableStd {
		c.Cmd.Stdin = os.Stdin
		c.Cmd.Stdout = os.Stdout
		c.Cmd.Stderr = os.Stderr
	} else {
		c.Cmd.Stdin = conf.Stdin
		c.Cmd.Stdout = conf.Stdout
		c.Cmd.Stderr = conf.Stderr
	}
	sysProcAttr(c.Cmd)
	return c.Cmd
}

type Shell struct {
}

func (*Shell) init(args string) *exec.Cmd {
	var args_ []string
	var cmd *exec.Cmd
	switch runtime.GOOS {
	case "darwin":
		fallthrough
	case "linux":
		args_ = append(args_, "-c")
		args_ = append(args_, args)
		cmd = exec.Command(os.Getenv("SHELL"), args_...)
		break
	case "windows":
		args_ = append(args_, "/C")
		args_ = append(args_, args)
		cmd = exec.Command("cmd", args_...)
		break
	default:
		os.Exit(1)
	}
	return cmd
}

func (s *Shell) Start(args []string, config Config, env ...string) error {
	var c = NewCmd(s.init(strings.Join(args, " ")), &config, env...)
	return c.Start()
}

// Run 阻塞式同步执行
func (s *Shell) Run(cmd string) Result {
	res := Result{Command: cmd, StartTime: time.Now()}
	conf := Config{Stdout: &res.stdout, Stderr: &res.stderr}
	var c = NewCmd(s.init(cmd), &conf)
	res.Err = c.Cmd.Run()
	res.EndTime = time.Now()
	return res
}

// RunTimeout 阻塞式超时同步执行
func (s *Shell) RunTimeout(cmd string, second time.Duration) Result {
	res := Result{Command: cmd, StartTime: time.Now()}
	conf := Config{Stdout: &res.stdout, Stderr: &res.stderr}
	c := NewCmd(s.init(cmd), &conf)
	tt := second * time.Second
	if tt <= 0 {
		tt = 30 * time.Second
	}
	t := time.NewTimer(tt)
	defer t.Stop()
	stop := make(chan struct{}, 1)
	ChildRunning(func() {
		res.Err = c.Cmd.Run()
		res.EndTime = time.Now()
		stop <- struct{}{}
	})
	select {
	case <-stop:
		break

	case <-t.C:
		if c.Cmd.Process != nil {
			err := c.Cmd.Process.Kill()
			res.Err = fmt.Errorf("cmd time out, kill the process id= %d,%v", c.Cmd.Process.Pid, err)
			return res
		}
		res.Err = fmt.Errorf("cmd time out")
		return res

	case <-Closed():
		break
	}
	return res
}

type Cmd struct {
	Cmd       *exec.Cmd
	Err       error
	Data      interface{}
	asyncWait *RunningCheck
}

func NewArgCmd(id int, args []string, conf *Config, opts ...ArgOptionFunc) (*Cmd, error) {
	var opt = ArgOption{Id: id, Args: args}
	for i := range opts {
		opts[i](&opt)
	}
	if len(opt.AppExe) == 0 {
		opt.AppExe = opt.Args[0]
	}
	fileName, err := filepath.Abs(opt.AppExe)
	if err != nil {
		return nil, err
	}
	return NewCmd(exec.Command(fileName, opt.Args[1:]...), conf, opt.Env...).WithData(opt.Data), nil
}

func NewCmd(cmd *exec.Cmd, conf *Config, env ...string) *Cmd {
	var c = &Cmd{Cmd: cmd}
	c.init(conf, env...)
	return c
}

func (c *Cmd) WithErr(err error) *Cmd {
	c.Err = err
	return c
}

func (c *Cmd) WithData(da interface{}) *Cmd {
	c.Data = da
	return c
}

func (c *Cmd) Start() error {
	if c.Cmd == nil {
		return errors.New("cmd not exist")
	}
	return c.Cmd.Start()
}

func (c *Cmd) Kill() error {
	if c.Cmd == nil || c.Cmd.Process == nil {
		return errors.New("cmd not exist")
	}
	return c.Cmd.Process.Kill()
}

func (c *Cmd) Wait() error {
	if c == nil || c.Cmd == nil {
		return errors.New("cmd not exist")
	}
	return c.Cmd.Wait()
}

// AsyncWait 异步
func (c *Cmd) AsyncWait(q func(cmd *Cmd, err error)) {
	if c == nil || c.Cmd == nil {
		q(nil, errors.New("cmd not exist"))
		return
	}
	if c.asyncWait == nil {
		c.asyncWait = &RunningCheck{}
	}
	c.asyncWait.GoRunning(func() {
		err := c.Cmd.Wait()
		q(c, err)
	})
}

func StripArgs(args []string, arg string) []string {
	ll := len(args)
	for i := 0; i < ll; {
		if args[i] == arg {
			next := 1
			if i+1 < ll && args[i+1][0] != '-' {
				next = 2
			}
			args = append(args[:i], args[i+next:]...)
			break
		}
		i++
	}
	return args
}

// SyncExec 超时同步执行
// timeout = 0, default 30 seconds
func SyncExec(cmd string, second time.Duration) Result {
	var s = Shell{}
	return s.RunTimeout(cmd, second)
}

func SyncCmd(ctx context.Context, command, env []string, dir string) (Result, error) {
	res := Result{StartTime: time.Now()}
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	cmd.Env = env
	cmd.Dir = dir
	out, err := cmd.CombinedOutput()
	res.stdout.Write(out)
	res.EndTime = time.Now()
	if err != nil {
		res.Err = err
		return res, err
	}
	return res, nil
}

// AsyncExec 异步执行,返回chan
func AsyncExec(cmd string) <-chan Result {
	resCh := make(chan Result, 1)
	ChildRunning(func() {
		var s = Shell{}
		res := s.Run(cmd)
		resCh <- res
	})
	return resCh
}

// CreateShellFile 创建临时的 shell 脚本文件
// content 创建的脚本内容
func CreateShellFile(pattern, content string) (tmpFile string, err error) {
	file, err := os.CreateTemp("", pattern)
	if err != nil {
		return
	}
	defer func() {
		file.Close()
		switch runtime.GOOS {
		case "windows":
			tmpFile = file.Name() + ".bat"
		default:
			tmpFile = file.Name() + ".sh"
		}
		err = os.Rename(file.Name(), tmpFile)
	}()

	err = file.Chmod(0777)
	if err != nil {
		return "", err
	}
	switch runtime.GOOS {
	case "windows":
	default:
		_, err = file.WriteString("#!/bin/bash\n")
		if err != nil {
			return
		}
		_, err = file.WriteString("set -e\n")
		if err != nil {
			return
		}
	}
	_, err = file.WriteString(content)
	return
}