修正shell判定

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)

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

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

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
+}

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())
+}