Add set_env option

Signed-off-by: Mingxiang Xue <mingxiangxue@gmail.com>

6 files changed, 33 insertions, 3 deletions

diff --git a/lib/net/ssh.rb b/lib/net/ssh.rb
index 542c757..9df01f0 100644
--- a/lib/net/ssh.rb
+++ b/lib/net/ssh.rb
@@ -4,6 +4,7 @@ ENV['HOME'] ||= ENV['HOMEPATH'] ? "#{ENV['HOMEDRIVE']}#{ENV['HOMEPATH']}" : Dir.
 
 require 'logger'
 require 'etc'
+require 'shellwords'
 
 require 'net/ssh/config'
 require 'net/ssh/errors'
@@ -70,7 +71,7 @@ module Net
       rekey_blocks_limit rekey_limit rekey_packet_limit timeout verbose
       known_hosts global_known_hosts_file user_known_hosts_file host_key_alias
       host_name user properties passphrase keys_only max_pkt_size
-      max_win_size send_env use_agent number_of_password_prompts
+      max_win_size send_env set_env use_agent number_of_password_prompts
       append_all_supported_algorithms non_interactive password_prompt
       agent_socket_factory minimum_dh_bits verify_host_key
       fingerprint_hash check_host_ip
@@ -175,6 +176,8 @@ module Net
     # * :rekey_packet_limit => the max number of packets to process before rekeying
     # * :send_env => an array of local environment variable names to export to the
     #   remote environment. Names may be given as String or Regexp.
+    # * :set_env => a hash of environment variable names and values to set to the
+    #   remote environment. Override the ones if specified in +send_env+.
     # * :timeout => how long to wait for the initial connection to be made
     # * :user => the user name to log in as; this overrides the +user+
     #   parameter, and is primarily only useful when provided via an SSH
diff --git a/lib/net/ssh/config.rb b/lib/net/ssh/config.rb
index 5c81e61..a40262e 100644
--- a/lib/net/ssh/config.rb
+++ b/lib/net/ssh/config.rb
@@ -276,6 +276,8 @@ module Net
           when :sendenv
             multi_send_env = value.to_s.split(/\s+/)
             hash[:send_env] = multi_send_env.map { |e| Regexp.new pattern2regex(e).source, false }
+          when :setenv
+            hash[:set_env] = Shellwords.split(value.to_s).map { |e| e.split '=', 2 }.to_h
           when :numberofpasswordprompts
             hash[:number_of_password_prompts] = value.to_i
           when *TRANSLATE_CONFIG_KEY_RENAME_MAP.keys
diff --git a/lib/net/ssh/connection/channel.rb b/lib/net/ssh/connection/channel.rb
index a47d77d..7bfee72 100644
--- a/lib/net/ssh/connection/channel.rb
+++ b/lib/net/ssh/connection/channel.rb
@@ -2,8 +2,8 @@ require 'net/ssh/loggable'
 require 'net/ssh/connection/constants'
 require 'net/ssh/connection/term'
 
-module Net 
-  module SSH 
+module Net
+  module SSH
     module Connection
 
       # The channel abstraction. Multiple "channels" can be multiplexed onto a
@@ -530,6 +530,7 @@ module Net
           @remote_maximum_packet_size = max_packet
           connection.forward.agent(self) if connection.options[:forward_agent] && type == "session"
           forward_local_env(connection.options[:send_env]) if connection.options[:send_env]
+          set_remote_env(connection.options[:set_env]) if connection.options[:set_env]
           @on_confirm_open.call(self) if @on_confirm_open
         end
     
@@ -677,6 +678,13 @@ module Net
             end
           end
         end
+
+        # Set a +Hash+ of environment variables in the remote process' environment.
+        #
+        #   channel.set_remote_env foo: 'bar', baz: 'whale'
+        def set_remote_env(env)
+          env.each { |key, value| self.env(key, value) }
+        end
       end
 
     end
diff --git a/test/configs/set_env b/test/configs/set_env
new file mode 100644
index 0000000..35df6c0
--- /dev/null
+++ b/test/configs/set_env
@@ -0,0 +1,2 @@
+Host 1234
+	SetEnv foo="bar" baz=whale cat="black hole"
diff --git a/test/start/test_options.rb b/test/start/test_options.rb
index 26c2a14..43a88a7 100644
--- a/test/start/test_options.rb
+++ b/test/start/test_options.rb
@@ -32,6 +32,13 @@ module NetSSH
       end
     end
 
+    def test_start_should_accept_set_env_option
+      assert_nothing_raised do
+        options = { set_env: { foo: 'bar', baz: 'whale will' } }
+        Net::SSH.start('localhost', 'testuser', options)
+      end
+    end
+
     def test_start_should_accept_number_of_password_prompts_option
       assert_nothing_raised do
         options = { number_of_password_prompts: 2 }
diff --git a/test/test_config.rb b/test/test_config.rb
index eab6d7f..aa21ad6 100644
--- a/test/test_config.rb
+++ b/test/test_config.rb
@@ -138,6 +138,7 @@ class TestConfig < NetSSHTest
       'pubkeyauthentication'    => true,
       'rekeylimit'              => 1024,
       'sendenv'                 => "LC_*",
+      'setenv'                  => 'foo="bar" baz=whale cat="black hole"',
       'numberofpasswordprompts' => '123',
       'serveraliveinterval'     => '2',
       'serveralivecountmax'     => '4',
@@ -162,6 +163,7 @@ class TestConfig < NetSSHTest
     assert_equal 1024,      net_ssh[:rekey_limit]
     assert_equal "127.0.0.1", net_ssh[:bind_address]
     assert_equal [/^LC_.*$/], net_ssh[:send_env]
+    assert_equal Hash['foo'=>'bar', 'baz'=>'whale', 'cat'=>'black hole'], net_ssh[:set_env]
     assert_equal 123,       net_ssh[:number_of_password_prompts]
     assert_equal 4,         net_ssh[:keepalive_maxcount]
     assert_equal 2,         net_ssh[:keepalive_interval]
@@ -285,6 +287,12 @@ class TestConfig < NetSSHTest
     assert_equal [/^GIT_.*$/, /^LANG$/, /^LC_.*$/], net_ssh[:send_env]
   end
 
+  def test_load_with_set_env
+    config = Net::SSH::Config.load(config(:set_env), '1234')
+    net_ssh = Net::SSH::Config.translate(config)
+    assert_equal Hash['foo'=>'bar', 'baz'=>'whale', 'cat'=>'black hole'], net_ssh[:set_env]
+  end
+
   def test_load_with_remote_user
     config = Net::SSH::Config.load(config(:proxy_remote_user), "behind-proxy")
     net_ssh = Net::SSH::Config.translate(config)