From 577ecdb54f5e379bc427cb5ea8c563fed6c5eca0 Mon Sep 17 00:00:00 2001 From: Lonny Wong Date: Sat, 11 Nov 2023 16:33:09 +0800 Subject: [PATCH] support send and set env SendEnv SetEnv --- tssh/args.go | 22 +++++-- tssh/args_test.go | 65 +++++++++++++++++++ tssh/config.go | 8 +++ tssh/env.go | 157 ++++++++++++++++++++++++++++++++++++++++++++++ tssh/login.go | 7 ++- 5 files changed, 254 insertions(+), 5 deletions(-) create mode 100644 tssh/env.go diff --git a/tssh/args.go b/tssh/args.go index 24949f4..81ae6e1 100644 --- a/tssh/args.go +++ b/tssh/args.go @@ -26,6 +26,7 @@ SOFTWARE. package tssh import ( + "bytes" "fmt" "strings" ) @@ -86,15 +87,28 @@ func (sshArgs) Version() string { } func (o *sshOption) UnmarshalText(b []byte) error { - s := string(b) - pos := strings.Index(s, "=") - if pos < 1 || strings.TrimSpace(s[pos+1:]) == "" { + s := string(bytes.TrimSpace(b)) + pos := strings.IndexRune(s, '=') + if pos >= 0 { + p := strings.IndexAny(strings.TrimRight(s[:pos], " \t"), " \t") + if p > 0 { + pos = p + } + } else { + pos = strings.IndexAny(s, " \t") + } + if pos < 0 { + return fmt.Errorf("invalid option: %s", s) + } + key := strings.TrimSpace(s[:pos]) + value := strings.TrimSpace(s[pos+1:]) + if key == "" || value == "" { return fmt.Errorf("invalid option: %s", s) } if o.options == nil { o.options = make(map[string]string) } - o.options[strings.ToLower(strings.TrimSpace(s[:pos]))] = strings.TrimSpace(s[pos+1:]) + o.options[strings.ToLower(key)] = value return nil } diff --git a/tssh/args_test.go b/tssh/args_test.go index 64f05bd..a3cd3d2 100644 --- a/tssh/args_test.go +++ b/tssh/args_test.go @@ -131,3 +131,68 @@ func TestSshArgs(t *testing.T) { assertArgsError("-L", "missing value for -L") assertArgsError("-R", "missing value for -R") } + +func TestSshOption(t *testing.T) { + assert := assert.New(t) + assertRemoteCommand := func(optionArg, optionValue string) { + t.Helper() + var args sshArgs + p, err := arg.NewParser(arg.Config{}, &args) + assert.Nil(err) + err = p.Parse([]string{optionArg}) + assert.Nil(err) + assert.Equal(sshArgs{Option: sshOption{map[string]string{"remotecommand": optionValue}}}, args) + } + + assertRemoteCommand("-oRemoteCommand echo abc", "echo abc") + assertRemoteCommand("-o RemoteCommand echo abc", "echo abc") + assertRemoteCommand("-o\tRemoteCommand\techo\tabc", "echo\tabc") + + assertRemoteCommand("-oRemoteCommand echo = abc", "echo = abc") + assertRemoteCommand("-o RemoteCommand echo = abc ", "echo = abc") + assertRemoteCommand("-o\tRemoteCommand \techo \t= \tabc \t", "echo \t= \tabc") + + assertRemoteCommand("-oRemoteCommand=echo abc", "echo abc") + assertRemoteCommand("-o RemoteCommand = echo abc ", "echo abc") + assertRemoteCommand("-o\tRemoteCommand\t=\techo abc ", "echo abc") + + assertRemoteCommand("-oRemoteCommand = echo abc ", "echo abc") + assertRemoteCommand("-o RemoteCommand = echo abc ", "echo abc") + assertRemoteCommand("-o \tRemoteCommand \t= \techo \tabc\t ", "echo \tabc") + + assertRemoteCommand("-oRemoteCommand = echo = abc ", "echo = abc") + assertRemoteCommand("-o RemoteCommand = echo = abc ", "echo = abc") + assertRemoteCommand("-o \tRemoteCommand\t =\t echo\t =\t abc \t", "echo\t =\t abc") + + assertInvalidOption := func(optionArg string) { + t.Helper() + var args sshArgs + p, err := arg.NewParser(arg.Config{}, &args) + assert.Nil(err) + err = p.Parse([]string{optionArg}) + assert.NotNil(err) + if err != nil { + assert.Contains(err.Error(), "invalid option") + } + } + + assertInvalidOption("-oRemoteCommand") + assertInvalidOption("-oRemoteCommand ") + assertInvalidOption("-oRemoteCommand \t ") + assertInvalidOption("-oRemoteCommand=") + assertInvalidOption("-oRemoteCommand = ") + assertInvalidOption("-oRemoteCommand \t = \t ") + + assertInvalidOption("-o \t RemoteCommand") + assertInvalidOption("-o \t RemoteCommand ") + assertInvalidOption("-o \t RemoteCommand \t ") + assertInvalidOption("-o \t RemoteCommand=") + assertInvalidOption("-o \t RemoteCommand = ") + assertInvalidOption("-o \t RemoteCommand \t = \t ") + + assertInvalidOption("-o=RemoteCommand") + assertInvalidOption("-o =RemoteCommand") + assertInvalidOption("-o= RemoteCommand") + assertInvalidOption("-o = RemoteCommand") + assertInvalidOption("-o\t=\tRemoteCommand") +} diff --git a/tssh/config.go b/tssh/config.go index ce45ada..c8ce33e 100644 --- a/tssh/config.go +++ b/tssh/config.go @@ -359,6 +359,14 @@ func getOptionConfig(args *sshArgs, option string) string { return getConfig(args.Destination, option) } +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 +} + func getExOptionConfig(args *sshArgs, option string) string { if value := args.Option.get(option); value != "" { return value diff --git a/tssh/env.go b/tssh/env.go new file mode 100644 index 0000000..cf51841 --- /dev/null +++ b/tssh/env.go @@ -0,0 +1,157 @@ +/* +MIT License + +Copyright (c) 2023 Lonny Wong +Copyright (c) 2023 [Contributors](https://github.com/trzsz/trzsz-ssh/graphs/contributors) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package tssh + +import ( + "fmt" + "os" + "regexp" + "strings" + + "github.com/google/shlex" + "golang.org/x/crypto/ssh" +) + +type sshEnv struct { + name string + value string +} + +func getSendEnvs(args *sshArgs) ([]*sshEnv, error) { + envSet := make(map[string]struct{}) + for _, envCfg := range getAllOptionConfig(args, "SendEnv") { + for _, env := range strings.Fields(envCfg) { + if len(env) > 0 { + envSet[env] = struct{}{} + } + } + } + if len(envSet) == 0 { + return nil, nil + } + + var buf strings.Builder + for env := range envSet { + if buf.Len() > 0 { + buf.WriteRune('|') + } + buf.WriteString("(^") + for _, c := range env { + switch c { + case '*': + buf.WriteString(".*") + case '?': + buf.WriteRune('.') + case '(', ')', '[', ']', '{', '}', '.', '+', ',', '-', '^', '$', '|', '\\': + buf.WriteRune('\\') + buf.WriteRune(c) + default: + buf.WriteRune(c) + } + } + buf.WriteString("$)") + } + expr := buf.String() + debug("send env regexp: %s", expr) + + re, err := regexp.Compile(expr) + if err != nil { + return nil, fmt.Errorf("compile SendEnv regexp failed: %v", err) + } + + var envs []*sshEnv + for _, env := range os.Environ() { + var name string + pos := strings.IndexRune(env, '=') + if pos < 0 { + name = strings.TrimSpace(env) + } else { + name = strings.TrimSpace(env[:pos]) + } + if !re.MatchString(name) { + continue + } + var value string + if pos >= 0 { + value = strings.TrimSpace(env[pos+1:]) + } + envs = append(envs, &sshEnv{name, value}) + } + return envs, nil +} + +func getSetEnvs(args *sshArgs) ([]*sshEnv, error) { + envCfg := getOptionConfig(args, "SetEnv") + if envCfg == "" { + return nil, nil + } + tokens, err := shlex.Split(envCfg) + if err != nil { + return nil, fmt.Errorf("invalid SetEnv: %s", envCfg) + } + var envs []*sshEnv + for _, token := range tokens { + pos := strings.IndexRune(token, '=') + if pos < 0 { + return nil, fmt.Errorf("invalid SetEnv: %s", envCfg) + } + name := strings.TrimSpace(token[:pos]) + if name == "" { + return nil, fmt.Errorf("invalid SetEnv: %s", envCfg) + } + value := strings.TrimSpace(token[pos+1:]) + envs = append(envs, &sshEnv{name, value}) + } + return envs, nil +} + +func sendAndSetEnv(args *sshArgs, session *ssh.Session) error { + envs, err := getSendEnvs(args) + if err != nil { + return err + } + for _, env := range envs { + if err := session.Setenv(env.name, env.value); err != nil { + debug("send env failed: %s = \"%s\"", env.name, env.value) + } else { + debug("send env success: %s = \"%s\"", env.name, env.value) + } + } + + envs, err = getSetEnvs(args) + if err != nil { + return err + } + for _, env := range envs { + if err := session.Setenv(env.name, env.value); err != nil { + debug("set env failed: %s = \"%s\"", env.name, env.value) + } else { + debug("set env success: %s = \"%s\"", env.name, env.value) + } + } + + return nil +} diff --git a/tssh/login.go b/tssh/login.go index a68eb48..1e4b9c2 100644 --- a/tssh/login.go +++ b/tssh/login.go @@ -895,7 +895,11 @@ func sshLogin(args *sshArgs, tty bool) (client *ssh.Client, session *ssh.Session err = fmt.Errorf("ssh new session failed: %v", err) return } - session.Stderr = os.Stderr + + // send and set env + if err = sendAndSetEnv(args, session); err != nil { + return + } // session input and output serverIn, err = session.StdinPipe() @@ -908,6 +912,7 @@ func sshLogin(args *sshArgs, tty bool) (client *ssh.Client, session *ssh.Session err = fmt.Errorf("stdout pipe failed: %v", err) return } + session.Stderr = os.Stderr // ssh agent forward sshAgentForward(args, client, session)