From ce7a67f4d0f4b7f682295f8c191d22850da6fd8f Mon Sep 17 00:00:00 2001 From: 兔子 Date: Mon, 8 Jun 2020 15:11:43 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=AD=A3shell=E5=88=A4=E5=AE=9A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crypto.go | 6 +- shell.go | 7 +- starshell.go | 161 +++++++++++++++++++++++++++++++++++++++++++++ starshell_linux.go | 34 ++++++++++ starshell_test.go | 60 +++++++++++++++++ 5 files changed, 260 insertions(+), 8 deletions(-) create mode 100644 starshell.go create mode 100644 starshell_linux.go create mode 100644 starshell_test.go diff --git a/crypto.go b/crypto.go index ee9f008..686633c 100644 --- a/crypto.go +++ b/crypto.go @@ -823,17 +823,17 @@ func FillWithRandom(filepath string, filesize int, bufcap int, bufnum int, shell for i := 0; i < bufnum; i++ { buftmp = []byte{} for j := 0; j < bufcap; j++ { - buftmp = append(buftmp, byte(rands.Intn(255))) + buftmp = append(buftmp, byte(rands.Intn(256))) } buf = append(buf, buftmp) } sum := 0 for { if filesize-sum < bufcap { - writer.Write(buf[rands.Intn(bufnum-1)][0 : filesize-sum]) + writer.Write(buf[rands.Intn(bufnum)][0 : filesize-sum]) sum += filesize - sum } else { - writer.Write(buf[rands.Intn(bufnum-1)]) + writer.Write(buf[rands.Intn(bufnum)]) sum += bufcap } go shell(float64(sum) / float64(filesize) * 100) diff --git a/shell.go b/shell.go index 74e0966..5c0305f 100644 --- a/shell.go +++ b/shell.go @@ -31,10 +31,7 @@ func NewPipeShell(command string, arg ...string) (*suncli, error) { lovecli.counter = 0 cmd := exec.Command(command, arg...) lovecli.cmd = cmd - lovecli.infile, err = lovecli.cmd.StdinPipe() - if err != nil { - return &lovecli, err - } + lovecli.outfile, err = lovecli.cmd.StdoutPipe() if err != nil { return &lovecli, err @@ -119,7 +116,7 @@ func (this suncli) WriteCmd(cmdstr string) { return } -func (this suncli) ExitCode() int { +func (this *suncli) ExitCode() int { return this.cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus() } diff --git a/starshell.go b/starshell.go new file mode 100644 index 0000000..635fdef --- /dev/null +++ b/starshell.go @@ -0,0 +1,161 @@ +package starainrt + +import ( + "context" + "errors" + "io" + "os" + "os/exec" + "syscall" + "time" +) + +type StarShell struct { + outfile *io.ReadCloser + infile *io.WriteCloser + errfile *io.ReadCloser + cmd *exec.Cmd + running bool + stopSign context.Context + stopFunc context.CancelFunc + stdout []byte + errout []byte + runerr error + exitcode int +} + +func NewStarShell(command string, args ...string) (*StarShell, error) { + shell := new(StarShell) + shell.stopSign, shell.stopFunc = context.WithCancel(context.Background()) + cmd := exec.CommandContext(shell.stopSign, command, args...) + shell.cmd = cmd + infile, err := shell.cmd.StdinPipe() + if err != nil { + return shell, err + } + shell.infile = &infile + errfile, err := shell.cmd.StderrPipe() + if err != nil { + return shell, err + } + shell.errfile = &errfile + outfile, err := shell.cmd.StdoutPipe() + if err != nil { + return shell, err + } + shell.outfile = &outfile + shell.runerr = nil + shell.exitcode = -999 + return shell, nil +} + +func (starcli *StarShell) Start() error { + if err := starcli.cmd.Start(); err != nil { + return err + } + starcli.running = true + go func() { + err := starcli.cmd.Wait() + if err != nil { + starcli.runerr = err + } + starcli.exitcode = starcli.cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus() + starcli.running = false + }() + go starcli.queryResult() + go starcli.queryErrResult() + return nil +} + +func (starcli *StarShell) queryResult() error { + for starcli.running { + out := make([]byte, 65535) + n, err := (*starcli.outfile).Read(out) + if n != 0 { + for i := 0; i < n; i++ { + starcli.stdout = append(starcli.stdout, out[i]) + } + } + if err != nil { + if err == io.EOF { + break + } else { + starcli.runerr = err + return err + } + + } + time.Sleep(time.Microsecond * 100) + } + return nil +} + +func (starcli *StarShell) queryErrResult() error { + for starcli.running { + out := make([]byte, 65535) + n, err := (*starcli.errfile).Read(out) + if n != 0 { + for i := 0; i < n; i++ { + starcli.errout = append(starcli.errout, out[i]) + } + } + if err != nil { + if err == io.EOF { + break + } else { + starcli.runerr = err + return err + } + } + time.Sleep(time.Microsecond * 100) + } + return nil +} + +func (starcli *StarShell) GetResult() (string, error) { + np := len(starcli.stdout) + res1 := string(starcli.stdout[0:np]) + starcli.stdout = starcli.stdout[np:] + np = len(starcli.errout) + res2 := string(starcli.errout[0:np]) + starcli.errout = starcli.errout[np:] + if len(res2) == 0 && starcli.runerr != nil { + res2 = starcli.runerr.Error() + } + if len(res2) == 0 { + return res1, nil + } + return res1, errors.New(res2) +} + +func (starcli *StarShell) Exec(cmd string, wait int) (string, error) { + (*starcli.infile).Write([]byte(cmd + "\n")) + time.Sleep(time.Millisecond * time.Duration(wait)) + return starcli.GetResult() +} + +func (starcli *StarShell) WriteCmd(cmdstr string) { + (*starcli.infile).Write([]byte(cmdstr + "\n")) + return +} + +func (starcli *StarShell) ExitCode() int { + return starcli.exitcode +} + +func (starcli *StarShell) Kill() { + starcli.stopFunc() + starcli.running = false +} + +func (starcli *StarShell) IsRunning() bool { + return starcli.running +} + +func (starcli *StarShell) GetPid() int { + return starcli.cmd.Process.Pid +} + +func (starcli *StarShell) Signal(sig os.Signal) error { + return starcli.cmd.Process.Signal(sig) +} diff --git a/starshell_linux.go b/starshell_linux.go new file mode 100644 index 0000000..5aeba9f --- /dev/null +++ b/starshell_linux.go @@ -0,0 +1,34 @@ +// +build linux + +package starainrt + +import "syscall" + +func (starcli *StarShell) SetRunUser(uid, gid uint32) { + starcli.cmd.SysProcAttr = &syscall.SysProcAttr{ + Credential: &syscall.Credential{ + Uid: uid, + Gid: gid, + }, + Setsid: true, + } +} + +func (starcli *StarShell) Release() error { + if err := starcli.cmd.Start(); err != nil { + return err + } + starcli.cmd.SysProcAttr = &syscall.SysProcAttr{ + Setsid: true, + } + starcli.cmd.Process.Release() + return nil +} + +func (starcli *StarShell) SetKeepCaps() error { + _, _, err := syscall.RawSyscall(157 /*SYS PRCTL */, 0x8 /*PR SET KEEPCAPS*/, 1, 0) + if 0 != err { + return err + } + return nil +} diff --git a/starshell_test.go b/starshell_test.go new file mode 100644 index 0000000..c337612 --- /dev/null +++ b/starshell_test.go @@ -0,0 +1,60 @@ +package starainrt + +import ( + "fmt" + "testing" + "time" +) + +func Test_Starshell(t *testing.T) { + star, err := NewStarShell("cmd.exe") + if err != nil { + fmt.Println(err) + return + } + go func() { + time.Sleep(time.Second * 2) + star.WriteCmd("chcp 65001") + time.Sleep(time.Second * 2) + star.WriteCmd("ping baidu.com -n 10") + time.Sleep(time.Second * 12) + star.WriteCmd("exit") + }() + for star.IsRunning() { + time.Sleep(time.Millisecond * 100) + str, err := star.GetResult() + if err != nil { + fmt.Println("error:", err) + } + fmt.Print(str) + } + fmt.Println(star.ExitCode()) +} + +func Test_Starbash(t *testing.T) { + star, err := NewStarShell("bash", "-c", "ping baidu.com -c 10") + err = star.Start() + if err != nil { + fmt.Println(err) + return + } + /* + go func() { + time.Sleep(time.Second * 2) + star.WriteCmd("chcp 65001") + time.Sleep(time.Second * 2) + star.WriteCmd("ping baidu.com -n 10") + time.Sleep(time.Second * 12) + star.WriteCmd("exit") + }() + */ + for star.IsRunning() { + time.Sleep(time.Millisecond * 100) + str, err := star.GetResult() + if err != nil { + fmt.Println("error:", err) + } + fmt.Print(str) + } + fmt.Println(star.ExitCode()) +}