Skip to content

Commit

Permalink
support multiple options
Browse files Browse the repository at this point in the history
  • Loading branch information
@lonnywong
lonnywong committed Nov 12, 2023
1 parent faef20f commit f7e3927
Show file tree
Hide file tree
Showing 8 changed files with 284 additions and 224 deletions.
17 changes: 14 additions & 3 deletions tssh/args.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ import (
)

type sshOption struct {
options map[string]string
options map[string][]string
}

type multiStr struct {
Expand Down Expand Up @@ -106,16 +106,27 @@ func (o *sshOption) UnmarshalText(b []byte) error {
return fmt.Errorf("invalid option: %s", s)
}
if o.options == nil {
o.options = make(map[string]string)
o.options = make(map[string][]string)
}
o.options[strings.ToLower(key)] = value
o.options[strings.ToLower(key)] = append(o.options[strings.ToLower(key)], value)
return nil
}

func (o *sshOption) get(option string) string {
if o.options == nil {
return ""
}
values := o.options[strings.ToLower(option)]
if len(values) == 0 {
return ""
}
return values[0]
}

func (o *sshOption) getAll(option string) []string {
if o.options == nil {
return nil
}
return o.options[strings.ToLower(option)]
}

Expand Down
102 changes: 73 additions & 29 deletions tssh/args_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,33 +72,7 @@ func TestSshArgs(t *testing.T) {
assertArgsEqual("-Jjump", sshArgs{ProxyJump: "jump"})
assertArgsEqual("-J abc,def", sshArgs{ProxyJump: "abc,def"})
assertArgsEqual("-o RemoteCommand=none -oServerAliveInterval=5",
sshArgs{Option: sshOption{map[string]string{"remotecommand": "none", "serveraliveinterval": "5"}}})

newBindCfg := func(addr string, port int) *bindCfg {
return &bindCfg{&addr, port}
}
assertArgsEqual("-D 8000", sshArgs{DynamicForward: bindArgs{[]*bindCfg{{nil, 8000}}}})
assertArgsEqual("-D 127.0.0.1:8002", sshArgs{DynamicForward: bindArgs{[]*bindCfg{newBindCfg("127.0.0.1", 8002)}}})
assertArgsEqual("-D [fe80::6358:bbae:26f8:7859]:8003",
sshArgs{DynamicForward: bindArgs{[]*bindCfg{newBindCfg("fe80::6358:bbae:26f8:7859", 8003)}}})
assertArgsEqual("-D :8004 -D *:8005 -D ::1/8006",
sshArgs{DynamicForward: bindArgs{[]*bindCfg{newBindCfg("", 8004), newBindCfg("*", 8005), newBindCfg("::1", 8006)}}})

newForwardCfg := func(bindAddr string, bindPort int, destHost string, destPort int) *forwardCfg {
return &forwardCfg{&bindAddr, bindPort, destHost, destPort}
}
assertArgsEqual("-L 127.0.0.1:8001:[::1]:9001",
sshArgs{LocalForward: forwardArgs{[]*forwardCfg{newForwardCfg("127.0.0.1", 8001, "::1", 9001)}}})
assertArgsEqual("-L ::1/8002/localhost/9002",
sshArgs{LocalForward: forwardArgs{[]*forwardCfg{newForwardCfg("::1", 8002, "localhost", 9002)}}})
assertArgsEqual("-L 8003:0.0.0.0:9003 -L ::/8004/::1/9004", sshArgs{LocalForward: forwardArgs{
[]*forwardCfg{{nil, 8003, "0.0.0.0", 9003}, newForwardCfg("::", 8004, "::1", 9004)}}})
assertArgsEqual("-R :8001:[fe80::6358:bbae:26f8:7859]:9001",
sshArgs{RemoteForward: forwardArgs{[]*forwardCfg{newForwardCfg("", 8001, "fe80::6358:bbae:26f8:7859", 9001)}}})
assertArgsEqual("-R /8002/127.0.0.1/9002",
sshArgs{RemoteForward: forwardArgs{[]*forwardCfg{newForwardCfg("", 8002, "127.0.0.1", 9002)}}})
assertArgsEqual("-R 8003/::1/9003 -R *:8004:[fe80::6358:bbae:26f8:7859]:9004", sshArgs{RemoteForward: forwardArgs{
[]*forwardCfg{{nil, 8003, "::1", 9003}, newForwardCfg("*", 8004, "fe80::6358:bbae:26f8:7859", 9004)}}})
sshArgs{Option: sshOption{map[string][]string{"remotecommand": {"none"}, "serveraliveinterval": {"5"}}}})

assertArgsEqual("--reconnect", sshArgs{Reconnect: true})
assertArgsEqual("--dragfile", sshArgs{DragFile: true})
Expand All @@ -114,7 +88,7 @@ func TestSshArgs(t *testing.T) {

assertArgsEqual("-tp222 -oRemoteCommand=none -i~/.ssh/id_rsa -o ServerAliveCountMax=2 dest cmd arg1 arg2",
sshArgs{ForceTTY: true, Port: 222, Identity: multiStr{values: []string{"~/.ssh/id_rsa"}},
Option: sshOption{map[string]string{"remotecommand": "none", "serveralivecountmax": "2"}},
Option: sshOption{map[string][]string{"remotecommand": {"none"}, "serveralivecountmax": {"2"}}},
Destination: "dest", Command: "cmd", Argument: []string{"arg1", "arg2"}})

assertArgsError := func(cmdline, errMsg string) {
Expand All @@ -132,6 +106,59 @@ func TestSshArgs(t *testing.T) {
assertArgsError("-R", "missing value for -R")
}

func TestForwardArgs(t *testing.T) {
assert := assert.New(t)
assertDynamicForwardNil := func(argument string, address *string, port int) {
t.Helper()
var args sshArgs
p, err := arg.NewParser(arg.Config{}, &args)
assert.Nil(err)
err = p.Parse([]string{"-D", argument})
assert.Nil(err)
assert.Equal(sshArgs{DynamicForward: bindArgs{[]*bindCfg{{argument, address, port}}}}, args)
}
assertDynamicForward := func(argument string, address string, port int) {
t.Helper()
assertDynamicForwardNil(argument, &address, port)
}

assertDynamicForwardNil("8000", nil, 8000)
assertDynamicForward("127.0.0.1:8002", "127.0.0.1", 8002)
assertDynamicForward("[fe80::6358:bbae:26f8:7859]:8003", "fe80::6358:bbae:26f8:7859", 8003)
assertDynamicForward(":8004", "", 8004)
assertDynamicForward("*:8005", "*", 8005)
assertDynamicForward("::1/8006", "::1", 8006)

assertLRFwd := func(ftype, argument string, expectedArg sshArgs) {
t.Helper()
var args sshArgs
p, err := arg.NewParser(arg.Config{}, &args)
assert.Nil(err)
err = p.Parse([]string{ftype, argument})
assert.Nil(err)
assert.Equal(expectedArg, args)
}
assertLRForwardNil := func(argument string, bindAddr *string, bindPort int, destHost string, destPort int) {
t.Helper()
assertLRFwd("-L", argument, sshArgs{LocalForward: forwardArgs{[]*forwardCfg{
{argument, bindAddr, bindPort, destHost, destPort}}}})
assertLRFwd("-R", argument, sshArgs{RemoteForward: forwardArgs{[]*forwardCfg{
{argument, bindAddr, bindPort, destHost, destPort}}}})
}
assertLRForward := func(argument string, bindAddr string, bindPort int, destHost string, destPort int) {
t.Helper()
assertLRForwardNil(argument, &bindAddr, bindPort, destHost, destPort)
}
assertLRForward("127.0.0.1:8001:[::1]:9001", "127.0.0.1", 8001, "::1", 9001)
assertLRForward("::1/8002/localhost/9002", "::1", 8002, "localhost", 9002)
assertLRForwardNil("8003:0.0.0.0:9003", nil, 8003, "0.0.0.0", 9003)
assertLRForward("::/8004/::1/9004", "::", 8004, "::1", 9004)
assertLRForward(":8001:[fe80::6358:bbae:26f8:7859]:9001", "", 8001, "fe80::6358:bbae:26f8:7859", 9001)
assertLRForward("/8002/127.0.0.1/9002", "", 8002, "127.0.0.1", 9002)
assertLRForwardNil("8003/::1/9003", nil, 8003, "::1", 9003)
assertLRForward("*:8004:[fe80::6358:bbae:26f8:7859]:9004", "*", 8004, "fe80::6358:bbae:26f8:7859", 9004)
}

func TestSshOption(t *testing.T) {
assert := assert.New(t)
assertRemoteCommand := func(optionArg, optionValue string) {
Expand All @@ -141,7 +168,7 @@ func TestSshOption(t *testing.T) {
assert.Nil(err)
err = p.Parse([]string{optionArg})
assert.Nil(err)
assert.Equal(sshArgs{Option: sshOption{map[string]string{"remotecommand": optionValue}}}, args)
assert.Equal(sshArgs{Option: sshOption{map[string][]string{"remotecommand": {optionValue}}}}, args)
}

assertRemoteCommand("-oRemoteCommand echo abc", "echo abc")
Expand Down Expand Up @@ -196,3 +223,20 @@ func TestSshOption(t *testing.T) {
assertInvalidOption("-o = RemoteCommand")
assertInvalidOption("-o\t=\tRemoteCommand")
}

func TestMultiOptions(t *testing.T) {
assert := assert.New(t)
assertSendEnvs := func(optionArgs []string, optionValues ...string) {
t.Helper()
var args sshArgs
p, err := arg.NewParser(arg.Config{}, &args)
assert.Nil(err)
err = p.Parse(optionArgs)
assert.Nil(err)
assert.Equal(sshArgs{Option: sshOption{map[string][]string{"sendenv": optionValues}}}, args)
}

assertSendEnvs([]string{"-oSendEnv=ABC"}, "ABC")
assertSendEnvs([]string{"-oSendEnv=ABC 123", "-o", "SendEnv XYZ"}, "ABC 123", "XYZ")
assertSendEnvs([]string{"-o", "SendEnv ABC 123", "-oSendEnv = XYZ", "-oSendEnv m3"}, "ABC 123", "XYZ", "m3")
}
6 changes: 1 addition & 5 deletions tssh/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,11 +360,7 @@ func getOptionConfig(args *sshArgs, option string) string {
}

func getAllOptionConfig(args *sshArgs, option string) []string {
values := getAllConfig(args.Destination, option)
if value := args.Option.get(option); value != "" {
values = append(values, value)
}
return values
return append(args.Option.getAll(option), getAllConfig(args.Destination, option)...)
}

func getExOptionConfig(args *sshArgs, option string) string {
Expand Down
37 changes: 24 additions & 13 deletions tssh/ctrl_unix.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,8 @@ func startControlMaster(args *sshArgs) {
return
}

cmdArgs := []string{"-T", "-oClearAllForwardings=yes", "-oRemoteCommand=none", "-oConnectTimeout=5"}
cmdArgs := []string{"-T", "-oRemoteCommand=none", "-oConnectTimeout=5"}

if args.Debug {
cmdArgs = append(cmdArgs, "-v")
}
Expand All @@ -204,26 +205,36 @@ func startControlMaster(args *sshArgs) {
if args.Port != 0 {
cmdArgs = append(cmdArgs, "-p", strconv.Itoa(args.Port))
}
for _, identity := range args.Identity.values {
cmdArgs = append(cmdArgs, "-i", identity)
}
if args.ConfigFile != "" {
cmdArgs = append(cmdArgs, "-F", args.ConfigFile)
}
if args.ProxyJump != "" {
cmdArgs = append(cmdArgs, "-J", args.ProxyJump)
}

for key, value := range args.Option.options {
for _, identity := range args.Identity.values {
cmdArgs = append(cmdArgs, "-i", identity)
}
for _, b := range args.DynamicForward.binds {
cmdArgs = append(cmdArgs, "-D", b.argument)
}
for _, f := range args.LocalForward.cfgs {
cmdArgs = append(cmdArgs, "-L", f.argument)
}
for _, f := range args.RemoteForward.cfgs {
cmdArgs = append(cmdArgs, "-R", f.argument)
}

for key, values := range args.Option.options {
switch key {
case "controlmaster":
cmdArgs = append(cmdArgs, fmt.Sprintf("-oControlMaster=%s", value))
case "controlpath":
cmdArgs = append(cmdArgs, fmt.Sprintf("-oControlPath=%s", value))
case "controlpersist":
cmdArgs = append(cmdArgs, fmt.Sprintf("-oControlPersist=%s", value))
case "forwardagent":
cmdArgs = append(cmdArgs, fmt.Sprintf("-oForwardAgent=%s", value))
case "remotecommand":
break
case "enabletrzsz", "enabledragfile":
break
default:
for _, value := range values {
cmdArgs = append(cmdArgs, fmt.Sprintf("-o%s=%s", key, value))
}
}
}

Expand Down
Loading

0 comments on commit f7e3927

Please sign in to comment.