Merge pull request #149 from jonathanmorley/login-shell

Login shell

.rubocop.yml

@@ -69,3 +69,5 @@ Style/SpaceAroundOperators:
   Enabled: false
 Style/IfUnlessModifier:
   Enabled: false
+Style/FrozenStringLiteralComment:
+  Enabled: false

lib/train/extras/command_wrapper.rb

@@ -28,6 +28,9 @@ def run(_command)
   class LinuxCommand < CommandWrapperBase
     Train::Options.attach(self)
 
+    option :shell, default: false
+    option :shell_options, default: nil
+    option :shell_command, default: nil
     option :sudo, default: false
     option :sudo_options, default: nil
     option :sudo_password, default: nil
@@ -38,12 +41,14 @@ def initialize(backend, options)
       @backend = backend
       validate_options(options)
 
+      @shell = options[:shell]
+      @shell_options = options[:shell_options] # e.g. '--login'
+      @shell_command = options[:shell_command] # e.g. '/bin/sh'
       @sudo = options[:sudo]
       @sudo_options = options[:sudo_options]
       @sudo_password = options[:sudo_password]
       @sudo_command = options[:sudo_command]
       @user = options[:user]
-      @prefix = build_prefix
     end
 
     # (see CommandWrapperBase::verify)
@@ -71,29 +76,48 @@ def verify
 
     # (see CommandWrapperBase::run)
     def run(command)
-      @prefix + command
+      shell_wrap(sudo_wrap(command))
     end
 
     def self.active?(options)
-      options.is_a?(Hash) && options[:sudo]
+      options.is_a?(Hash) && (
+        options[:sudo] ||
+        options[:shell]
+      )
     end
 
     private
 
-    def build_prefix
-      return '' unless @sudo
-      return '' if @user == 'root'
+    # wrap the cmd in a sudo command
+    def sudo_wrap(cmd)
+      return cmd unless @sudo
+      return cmd if @user == 'root'
 
       res = (@sudo_command || 'sudo') + ' '
 
-      unless @sudo_password.nil?
-        b64pw = Base64.strict_encode64(@sudo_password + "\n")
-        res = "echo #{b64pw} | base64 --decode | #{res}-S "
-      end
+      res = "#{safe_string(@sudo_password + "\n")} | #{res}-S " unless @sudo_password.nil?
 
       res << @sudo_options.to_s + ' ' unless @sudo_options.nil?
 
-      res
+      res + cmd
+    end
+
+    # wrap the cmd in a subshell allowing for options to
+    # passed to the subshell
+    def shell_wrap(cmd)
+      return cmd unless @shell
+
+      shell = @shell_command || '$SHELL'
+      options = ' ' + @shell_options.to_s unless @shell_options.nil?
+
+      "#{safe_string(cmd)} | #{shell}#{options}"
+    end
+
+    # encapsulates encoding the string into a safe form, and decoding for use.
+    # @return [String] A command line snippet that can be used as part of a pipeline.
+    def safe_string(str)
+      b64str = Base64.strict_encode64(str)
+      "echo #{b64str} | base64 --decode"
     end
   end
 

test/unit/extras/command_wrapper_test.rb

@@ -16,47 +16,95 @@
     backend
   }
 
-  it 'wraps commands in sudo' do
-    lc = cls.new(backend, { sudo: true })
-    lc.run(cmd).must_equal "sudo #{cmd}"
-  end
+  describe 'sudo wrapping' do
+    it 'wraps commands in sudo' do
+      lc = cls.new(backend, { sudo: true })
+      lc.run(cmd).must_equal "sudo #{cmd}"
+    end
 
-  it 'doesnt wrap commands in sudo if user == root' do
-    lc = cls.new(backend, { sudo: true, user: 'root' })
-    lc.run(cmd).must_equal cmd
-  end
+    it 'doesnt wrap commands in sudo if user == root' do
+      lc = cls.new(backend, { sudo: true, user: 'root' })
+      lc.run(cmd).must_equal cmd
+    end
 
-  it 'wraps commands in sudo with all options' do
-    opts = rand.to_s
-    lc = cls.new(backend, { sudo: true, sudo_options: opts })
-    lc.run(cmd).must_equal "sudo #{opts} #{cmd}"
-  end
+    it 'wraps commands in sudo with all options' do
+      opts = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_options: opts })
+      lc.run(cmd).must_equal "sudo #{opts} #{cmd}"
+    end
 
-  it 'runs commands in sudo with password' do
-    pw = rand.to_s
-    lc = cls.new(backend, { sudo: true, sudo_password: pw })
-    bpw = Base64.strict_encode64(pw + "\n")
-    lc.run(cmd).must_equal "echo #{bpw} | base64 --decode | sudo -S #{cmd}"
-  end
+    it 'runs commands in sudo with password' do
+      pw = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_password: pw })
+      bpw = Base64.strict_encode64(pw + "\n")
+      lc.run(cmd).must_equal "echo #{bpw} | base64 --decode | sudo -S #{cmd}"
+    end
 
-  it 'wraps commands in sudo_command instead of sudo' do
-    sudo_command = rand.to_s
-    lc = cls.new(backend, { sudo: true, sudo_command: sudo_command })
-    lc.run(cmd).must_equal "#{sudo_command} #{cmd}"
-  end
+    it 'wraps commands in sudo_command instead of sudo' do
+      sudo_command = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_command: sudo_command })
+      lc.run(cmd).must_equal "#{sudo_command} #{cmd}"
+    end
+
+    it 'wraps commands in sudo_command with all options' do
+      opts = rand.to_s
+      sudo_command = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_command: sudo_command, sudo_options: opts })
+      lc.run(cmd).must_equal "#{sudo_command} #{opts} #{cmd}"
+    end
 
-  it 'wraps commands in sudo_command with all options' do
-    opts = rand.to_s
-    sudo_command = rand.to_s
-    lc = cls.new(backend, { sudo: true, sudo_command: sudo_command, sudo_options: opts })
-    lc.run(cmd).must_equal "#{sudo_command} #{opts} #{cmd}"
+    it 'runs commands in sudo_command with password' do
+      pw = rand.to_s
+      sudo_command = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_command: sudo_command, sudo_password: pw })
+      bpw = Base64.strict_encode64(pw + "\n")
+      lc.run(cmd).must_equal "echo #{bpw} | base64 --decode | #{sudo_command} -S #{cmd}"
+    end
   end
 
-  it 'runs commands in sudo_command with password' do
-    pw = rand.to_s
-    sudo_command = rand.to_s
-    lc = cls.new(backend, { sudo: true, sudo_command: sudo_command, sudo_password: pw })
-    bpw = Base64.strict_encode64(pw + "\n")
-    lc.run(cmd).must_equal "echo #{bpw} | base64 --decode | #{sudo_command} -S #{cmd}"
+  describe 'shell wrapping' do
+    it 'wraps commands in a default shell with login' do
+      lc = cls.new(backend, { shell: true, shell_options: '--login' })
+      bcmd = Base64.strict_encode64(cmd)
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | $SHELL --login"
+    end
+
+    it 'wraps sudo commands in a default shell with login' do
+      lc = cls.new(backend, { sudo: true, shell: true, shell_options: '--login' })
+      bcmd = Base64.strict_encode64("sudo #{cmd}")
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | $SHELL --login"
+    end
+
+    it 'wraps sudo commands and sudo passwords in a default shell with login' do
+      pw = rand.to_s
+      lc = cls.new(backend, { sudo: true, sudo_password: pw, shell: true, shell_options: '--login' })
+      bpw = Base64.strict_encode64(pw + "\n")
+      bcmd = Base64.strict_encode64("echo #{bpw} | base64 --decode | sudo -S #{cmd}")
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | $SHELL --login"
+      p bcmd
+    end
+
+    it 'wraps commands in a default shell when shell is true' do
+      lc = cls.new(backend, { shell: true })
+      bcmd = Base64.strict_encode64(cmd)
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | $SHELL"
+    end
+
+    it 'doesnt wrap commands in a shell when shell is false' do
+      lc = cls.new(backend, { shell: false })
+      lc.run(cmd).must_equal cmd
+    end
+
+    it 'wraps commands in a `shell` instead of default shell' do
+      lc = cls.new(backend, { shell: true, shell_command: '/bin/bash' })
+      bcmd = Base64.strict_encode64(cmd)
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | /bin/bash"
+    end
+
+    it 'wraps commands in a default shell with login' do
+      lc = cls.new(backend, { shell: true, shell_command: '/bin/bash', shell_options: '--login' })
+      bcmd = Base64.strict_encode64(cmd)
+      lc.run(cmd).must_equal "echo #{bcmd} | base64 --decode | /bin/bash --login"
+    end
   end
 end