github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/exec/cmd.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package executils
     4  
     5  import (
     6  	"bytes"
     7  	"context"
     8  	"os"
     9  	"os/exec"
    10  	"strings"
    11  	"time"
    12  )
    13  
    14  type Cmd struct {
    15  	name string
    16  	args []string
    17  	env  []string
    18  	dir  string
    19  
    20  	ctx        context.Context
    21  	timeout    time.Duration
    22  	cancelFunc func()
    23  
    24  	captureStdout bool
    25  	captureStderr bool
    26  
    27  	stdout *bytes.Buffer
    28  	stderr *bytes.Buffer
    29  
    30  	rawCmd *exec.Cmd
    31  }
    32  
    33  func NewCmd(name string, args ...string) *Cmd {
    34  	return &Cmd{
    35  		name: name,
    36  		args: args,
    37  	}
    38  }
    39  
    40  func NewTimeoutCmd(timeout time.Duration, name string, args ...string) *Cmd {
    41  	return (&Cmd{
    42  		name: name,
    43  		args: args,
    44  	}).WithTimeout(timeout)
    45  }
    46  
    47  func (this *Cmd) WithTimeout(timeout time.Duration) *Cmd {
    48  	this.timeout = timeout
    49  
    50  	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
    51  	this.ctx = ctx
    52  	this.cancelFunc = cancelFunc
    53  
    54  	return this
    55  }
    56  
    57  func (this *Cmd) WithStdout() *Cmd {
    58  	this.captureStdout = true
    59  	return this
    60  }
    61  
    62  func (this *Cmd) WithStderr() *Cmd {
    63  	this.captureStderr = true
    64  	return this
    65  }
    66  
    67  func (this *Cmd) WithEnv(env []string) *Cmd {
    68  	this.env = env
    69  	return this
    70  }
    71  
    72  func (this *Cmd) WithDir(dir string) *Cmd {
    73  	this.dir = dir
    74  	return this
    75  }
    76  
    77  func (this *Cmd) Start() error {
    78  	var cmd = this.compose()
    79  	return cmd.Start()
    80  }
    81  
    82  func (this *Cmd) Wait() error {
    83  	var cmd = this.compose()
    84  	return cmd.Wait()
    85  }
    86  
    87  func (this *Cmd) Run() error {
    88  	if this.cancelFunc != nil {
    89  		defer this.cancelFunc()
    90  	}
    91  
    92  	var cmd = this.compose()
    93  	return cmd.Run()
    94  }
    95  
    96  func (this *Cmd) RawStdout() string {
    97  	if this.stdout != nil {
    98  		return this.stdout.String()
    99  	}
   100  	return ""
   101  }
   102  
   103  func (this *Cmd) Stdout() string {
   104  	return strings.TrimSpace(this.RawStdout())
   105  }
   106  
   107  func (this *Cmd) RawStderr() string {
   108  	if this.stderr != nil {
   109  		return this.stderr.String()
   110  	}
   111  	return ""
   112  }
   113  
   114  func (this *Cmd) Stderr() string {
   115  	return strings.TrimSpace(this.RawStderr())
   116  }
   117  
   118  func (this *Cmd) String() string {
   119  	if this.rawCmd != nil {
   120  		return this.rawCmd.String()
   121  	}
   122  	var newCmd = exec.Command(this.name, this.args...)
   123  	return newCmd.String()
   124  }
   125  
   126  func (this *Cmd) Process() *os.Process {
   127  	if this.rawCmd != nil {
   128  		return this.rawCmd.Process
   129  	}
   130  	return nil
   131  }
   132  
   133  func (this *Cmd) compose() *exec.Cmd {
   134  	if this.rawCmd != nil {
   135  		return this.rawCmd
   136  	}
   137  
   138  	if this.ctx != nil {
   139  		this.rawCmd = exec.CommandContext(this.ctx, this.name, this.args...)
   140  	} else {
   141  		this.rawCmd = exec.Command(this.name, this.args...)
   142  	}
   143  
   144  	if this.env != nil {
   145  		this.rawCmd.Env = this.env
   146  	}
   147  
   148  	if len(this.dir) > 0 {
   149  		this.rawCmd.Dir = this.dir
   150  	}
   151  
   152  	if this.captureStdout {
   153  		this.stdout = &bytes.Buffer{}
   154  		this.rawCmd.Stdout = this.stdout
   155  	}
   156  	if this.captureStderr {
   157  		this.stderr = &bytes.Buffer{}
   158  		this.rawCmd.Stderr = this.stderr
   159  	}
   160  
   161  	return this.rawCmd
   162  }