add util

Author: wzshiming

util.go

@@ -0,0 +1,97 @@
+package commandproxy
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+)
+
+var (
+	errQuoteNotClose = fmt.Errorf("quote is not closed")
+	errEmptyCmmand   = fmt.Errorf("cmd is empty")
+)
+
+func SplitCommand(cmd string) ([]string, error) {
+	cmd = strings.TrimSpace(cmd)
+	if cmd == "" {
+		return nil, errEmptyCmmand
+	}
+	var ss []string
+
+	var buf bytes.Buffer
+	for len(cmd) != 0 {
+		i := strings.IndexAny(cmd, "'\"\\ \n\t")
+		if i == -1 {
+			buf.WriteString(cmd)
+			break
+		}
+		buf.WriteString(cmd[:i])
+		c := cmd[i]
+		cmd = cmd[i+1:]
+		switch c {
+		case '"':
+			i := strings.IndexByte(cmd, '"')
+			if i == -1 {
+				return nil, errQuoteNotClose
+			}
+			buf.WriteString(cmd[:i])
+			cmd = cmd[i+1:]
+		case '\'':
+			i := strings.IndexByte(cmd, '\'')
+			if i == -1 {
+				return nil, errQuoteNotClose
+			}
+			buf.WriteString(cmd[:i])
+			cmd = cmd[i+1:]
+		case '\\':
+			if len(cmd) == 0 {
+				break
+			}
+			if cmd[0] != '\n' {
+				buf.WriteByte(cmd[0])
+				cmd = cmd[:1]
+				continue
+			}
+			fallthrough
+		case '\t', ' ':
+			if buf.Len() != 0 {
+				ss = append(ss, buf.String())
+				buf.Reset()
+			}
+		}
+	}
+
+	if buf.Len() != 0 {
+		if len(ss) == 0 {
+			return []string{cmd}, nil
+		}
+		ss = append(ss, buf.String())
+	}
+
+	return ss, nil
+}
+
+func ReplaceEscape(s string, re map[byte]string) string {
+	re['%'] = "%"
+	var buf bytes.Buffer
+	for len(s) != 0 {
+		i := strings.IndexByte(s, '%')
+		if i == -1 || i >= len(s)-1 {
+			if buf.Len() == 0 {
+				return s
+			}
+			buf.WriteString(s)
+			break
+		}
+		buf.WriteString(s[:i])
+		s = s[i:]
+		rep, ok := re[s[1]]
+		if ok {
+			buf.WriteString(rep)
+		} else {
+			buf.WriteString(s[:2])
+		}
+		s = s[2:]
+	}
+	return buf.String()
+}

util_test.go

@@ -0,0 +1,215 @@
+package commandproxy
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestSplitCommand(t *testing.T) {
+	type args struct {
+		cmd string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    []string
+		wantErr bool
+	}{
+		{
+			args: args{
+				cmd: "",
+			},
+			wantErr: true,
+		},
+		{
+			args: args{
+				cmd: "a",
+			},
+			want: []string{"a"},
+		},
+		{
+			args: args{
+				cmd: "a b",
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: " a b",
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: "a b ",
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: "a  b",
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: "a  \"",
+			},
+			wantErr: true,
+		},
+		{
+			args: args{
+				cmd: "a  '",
+			},
+			wantErr: true,
+		},
+		{
+			args: args{
+				cmd: "a\t\tb",
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: `a \\`,
+			},
+			want: []string{"a", "\\"},
+		},
+		{
+			args: args{
+				cmd: `a \
+b`,
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: `a "b"`,
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: `a 'b'`,
+			},
+			want: []string{"a", "b"},
+		},
+		{
+			args: args{
+				cmd: `a "b "`,
+			},
+			want: []string{"a", "b "},
+		},
+		{
+			args: args{
+				cmd: `a """"`,
+			},
+			want: []string{"a"},
+		},
+
+		{
+			args: args{
+				cmd: `a "'b'"`,
+			},
+			want: []string{"a", "'b'"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := SplitCommand(tt.args.cmd)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("SplitCommand() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("SplitCommand() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestReplaceEscape(t *testing.T) {
+	type args struct {
+		s  string
+		re map[byte]string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			args: args{
+				s:  "",
+				re: map[byte]string{},
+			},
+			want: "",
+		},
+		{
+			args: args{
+				s:  "a",
+				re: map[byte]string{},
+			},
+			want: "a",
+		},
+		{
+			args: args{
+				s:  "%",
+				re: map[byte]string{},
+			},
+			want: "%",
+		},
+		{
+			args: args{
+				s:  "%%",
+				re: map[byte]string{},
+			},
+			want: "%",
+		},
+		{
+			args: args{
+				s:  "%%a",
+				re: map[byte]string{},
+			},
+			want: "%a",
+		},
+		{
+			args: args{
+				s:  "a%%",
+				re: map[byte]string{},
+			},
+			want: "a%",
+		},
+		{
+			args: args{
+				s:  "aa%%aa",
+				re: map[byte]string{},
+			},
+			want: "aa%aa",
+		},
+		{
+			args: args{
+				s:  "%a",
+				re: map[byte]string{},
+			},
+			want: "%a",
+		},
+		{
+			args: args{
+				s: "%h:%p",
+				re: map[byte]string{
+					'h': "127.0.0.1",
+					'p': "80",
+				},
+			},
+			want: "127.0.0.1:80",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := ReplaceEscape(tt.args.s, tt.args.re); got != tt.want {
+				t.Errorf("ReplaceEscape() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}