diff --git a/lib/train/errors.rb b/lib/train/errors.rb index b3fb8411..c3f70619 100644 --- a/lib/train/errors.rb +++ b/lib/train/errors.rb @@ -21,6 +21,9 @@ class ClientError < ::StandardError; end # in the transport layer. class TransportError < ::StandardError; end - # Exception for when no platform can be detected + # Exception for when no platform can be detected. class PlatformDetectionFailed < ::StandardError; end + + # Exception for when a invalid cache type is passed. + class UnknownCacheType < ::StandardError; end end diff --git a/lib/train/plugins/base_connection.rb b/lib/train/plugins/base_connection.rb index c2620071..77480806 100644 --- a/lib/train/plugins/base_connection.rb +++ b/lib/train/plugins/base_connection.rb @@ -15,9 +15,6 @@ class Train::Plugins::Transport class BaseConnection include Train::Extras - # Provide access to the files cache. - attr_reader :files - # Create a new Connection instance. # # @param options [Hash] connection options @@ -25,8 +22,35 @@ class BaseConnection def initialize(options = nil) @options = options || {} @logger = @options.delete(:logger) || Logger.new(STDOUT) - @files = {} Train::Platforms::Detect::Specifications::OS.load + + # default caching options + @cache_enabled = { + file: true, + command: false, + } + + @cache = {} + @cache_enabled.each_key do |type| + clear_cache(type) + end + end + + def cache_enabled?(type) + @cache_enabled[type.to_sym] + end + + # Enable caching types for Train. Currently we support + # :file and :command types + def enable_cache(type) + fail Train::UnknownCacheType, "#{type} is not a valid cache type" unless @cache_enabled.keys.include?(type.to_sym) + @cache_enabled[type.to_sym] = true + end + + def disable_cache(type) + fail Train::UnknownCacheType, "#{type} is not a valid cache type" unless @cache_enabled.keys.include?(type.to_sym) + @cache_enabled[type.to_sym] = false + clear_cache(type.to_sym) end # Closes the session connection, if it is still active. @@ -36,14 +60,14 @@ def close def to_json { - 'files' => Hash[@files.map { |x, y| [x, y.to_json] }], + 'files' => Hash[@cache[:file].map { |x, y| [x, y.to_json] }], } end def load_json(j) require 'train/transports/mock' j['files'].each do |path, jf| - @files[path] = Train::Transports::Mock::Connection::File.from_json(jf) + @cache[:file][path] = Train::Transports::Mock::Connection::File.from_json(jf) end end @@ -52,31 +76,29 @@ def local? false end - # Execute a command using this connection. - # - # @param command [String] command string to execute - # @return [CommandResult] contains the result of running the command - def run_command(_command) - fail Train::ClientError, "#{self.class} does not implement #run_command()" - end - # Get information on the operating system which this transport connects to. # # @return [Platform] system information def platform @platform ||= Train::Platforms::Detect.scan(self) end - # we need to keep os as a method for backwards compatibility with inspec alias os platform - # Interact with files on the target. Read, write, and get metadata - # from files via the transport. - # - # @param [String] path which is being inspected - # @return [FileCommon] file object that allows for interaction - def file(_path, *_args) - fail Train::ClientError, "#{self.class} does not implement #file(...)" + # This is the main command call for all connections. This will call the private + # run_command_via_connection on the connection with optional caching + def run_command(cmd) + return run_command_via_connection(cmd) unless cache_enabled?(:command) + + @cache[:command][cmd] ||= run_command_via_connection(cmd) + end + + # This is the main file call for all connections. This will call the private + # file_via_connection on the connection with optional caching + def file(path, *args) + return file_via_connection(path, *args) unless cache_enabled?(:file) + + @cache[:file][path] ||= file_via_connection(path, *args) end # Builds a LoginCommand which can be used to open an interactive @@ -84,7 +106,7 @@ def file(_path, *_args) # # @return [LoginCommand] array of command line tokens def login_command - fail Train::ClientError, "#{self.class} does not implement #login_command()" + fail NotImplementedError, "#{self.class} does not implement #login_command()" end # Block and return only when the remote host is prepared and ready to @@ -98,6 +120,27 @@ def wait_until_ready private + # Execute a command using this connection. + # + # @param command [String] command string to execute + # @return [CommandResult] contains the result of running the command + def run_command_via_connection(_command) + fail NotImplementedError, "#{self.class} does not implement #run_command_via_connection()" + end + + # Interact with files on the target. Read, write, and get metadata + # from files via the transport. + # + # @param [String] path which is being inspected + # @return [FileCommon] file object that allows for interaction + def file_via_connection(_path, *_args) + fail NotImplementedError, "#{self.class} does not implement #file_via_connection(...)" + end + + def clear_cache(type) + @cache[type.to_sym] = {} + end + # @return [Logger] logger for reporting information # @api private attr_reader :logger diff --git a/lib/train/transports/docker.rb b/lib/train/transports/docker.rb index 75ce59b8..6a893faf 100644 --- a/lib/train/transports/docker.rb +++ b/lib/train/transports/docker.rb @@ -69,18 +69,27 @@ def close # nothing to do at the moment end - def file(path) - @files[path] ||=\ - if os.aix? - Train::File::Remote::Aix.new(self, path) - elsif os.solaris? - Train::File::Remote::Unix.new(self, path) - else - Train::File::Remote::Linux.new(self, path) - end + def uri + if @container.nil? + "docker://#{@id}" + else + "docker://#{@container.id}" + end + end + + private + + def file_via_connection(path) + if os.aix? + Train::File::Remote::Aix.new(self, path) + elsif os.solaris? + Train::File::Remote::Unix.new(self, path) + else + Train::File::Remote::Linux.new(self, path) + end end - def run_command(cmd) + def run_command_via_connection(cmd) cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil? stdout, stderr, exit_status = @container.exec( [ @@ -93,13 +102,5 @@ def run_command(cmd) # @TODO: differentiate any other error raise end - - def uri - if @container.nil? - "docker://#{@id}" - else - "docker://#{@container.id}" - end - end end end diff --git a/lib/train/transports/local.rb b/lib/train/transports/local.rb index 1eeac343..b052524c 100644 --- a/lib/train/transports/local.rb +++ b/lib/train/transports/local.rb @@ -23,34 +23,35 @@ def initialize(options) @cmd_wrapper = CommandWrapper.load(self, options) end - def run_command(cmd) - cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil? - res = Mixlib::ShellOut.new(cmd) - res.run_command - CommandResult.new(res.stdout, res.stderr, res.exitstatus) - rescue Errno::ENOENT => _ - CommandResult.new('', '', 1) + def login_command + nil # none, open your shell + end + + def uri + 'local://' end def local? true end - def file(path) - @files[path] ||= \ - if os.windows? - Train::File::Local::Windows.new(self, path) - else - Train::File::Local::Unix.new(self, path) - end - end + private - def login_command - nil # none, open your shell + def run_command_via_connection(cmd) + cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil? + res = Mixlib::ShellOut.new(cmd) + res.run_command + CommandResult.new(res.stdout, res.stderr, res.exitstatus) + rescue Errno::ENOENT => _ + CommandResult.new('', '', 1) end - def uri - 'local://' + def file_via_connection(path) + if os.windows? + Train::File::Local::Windows.new(self, path) + else + Train::File::Local::Unix.new(self, path) + end end end end diff --git a/lib/train/transports/mock.rb b/lib/train/transports/mock.rb index 01f3a36e..b77e31ad 100644 --- a/lib/train/transports/mock.rb +++ b/lib/train/transports/mock.rb @@ -57,13 +57,13 @@ def trace_calls class Train::Transports::Mock class Connection < BaseConnection - attr_accessor :files, :commands attr_reader :os def initialize(conf = nil) super(conf) mock_os - @commands = {} + enable_cache(:file) + enable_cache(:command) end def uri @@ -90,8 +90,24 @@ def mock_os_hierarchy(plat) end end + def commands=(commands) + @cache[:command] = commands + end + + def commands + @cache[:command] + end + + def files=(files) + @cache[:file] = files + end + + def files + @cache[:file] + end + def mock_command(cmd, stdout = nil, stderr = nil, exit_status = 0) - @commands[cmd] = Command.new(stdout || '', stderr || '', exit_status) + @cache[:command][cmd] = Command.new(stdout || '', stderr || '', exit_status) end def command_not_found(cmd) @@ -104,24 +120,25 @@ def command_not_found(cmd) mock_command(cmd, nil, nil, 1) end - def run_command(cmd) - @commands[cmd] || - @commands[Digest::SHA256.hexdigest cmd.to_s] || - command_not_found(cmd) - end - def file_not_found(path) STDERR.puts('File not mocked: '+path.to_s) if @options[:verbose] File.new(self, path) end - def file(path) - @files[path] ||= file_not_found(path) - end - def to_s 'Mock Connection' end + + private + + def run_command_via_connection(cmd) + @cache[:command][Digest::SHA256.hexdigest cmd.to_s] || + command_not_found(cmd) + end + + def file_via_connection(path) + file_not_found(path) + end end end diff --git a/lib/train/transports/ssh_connection.rb b/lib/train/transports/ssh_connection.rb index 5e996796..6f02c147 100644 --- a/lib/train/transports/ssh_connection.rb +++ b/lib/train/transports/ssh_connection.rb @@ -55,63 +55,6 @@ def close @session = nil end - def file(path) - @files[path] ||= \ - if os.aix? - Train::File::Remote::Aix.new(self, path) - elsif os.solaris? - Train::File::Remote::Unix.new(self, path) - elsif os[:name] == 'qnx' - Train::File::Remote::Qnx.new(self, path) - else - Train::File::Remote::Linux.new(self, path) - end - end - - # (see Base::Connection#run_command) - def run_command(cmd) - stdout = stderr = '' - exit_status = nil - cmd.dup.force_encoding('binary') if cmd.respond_to?(:force_encoding) - logger.debug("[SSH] #{self} (#{cmd})") - - session.open_channel do |channel| - # wrap commands if that is configured - cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil? - - if @transport_options[:pty] - channel.request_pty do |_ch, success| - fail Train::Transports::SSHPTYFailed, 'Requesting PTY failed' unless success - end - end - - channel.exec(cmd) do |_, success| - abort 'Couldn\'t execute command on SSH.' unless success - - channel.on_data do |_, data| - stdout += data - end - - channel.on_extended_data do |_, _type, data| - stderr += data - end - - channel.on_request('exit-status') do |_, data| - exit_status = data.read_long - end - - channel.on_request('exit-signal') do |_, data| - exit_status = data.read_long - end - end - end - @session.loop - - CommandResult.new(stdout, stderr, exit_status) - rescue Net::SSH::Exception => ex - raise Train::Transports::SSHFailed, "SSH command failed (#{ex.message})" - end - # (see Base::Connection#login_command) def login_command level = logger.debug? ? 'VERBOSE' : 'ERROR' @@ -221,6 +164,61 @@ def establish_connection(opts) retry end + def file_via_connection(path) + if os.aix? + Train::File::Remote::Aix.new(self, path) + elsif os.solaris? + Train::File::Remote::Unix.new(self, path) + elsif os[:name] == 'qnx' + Train::File::Remote::Qnx.new(self, path) + else + Train::File::Remote::Linux.new(self, path) + end + end + + def run_command_via_connection(cmd) + stdout = stderr = '' + exit_status = nil + cmd.dup.force_encoding('binary') if cmd.respond_to?(:force_encoding) + logger.debug("[SSH] #{self} (#{cmd})") + + session.open_channel do |channel| + # wrap commands if that is configured + cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil? + + if @transport_options[:pty] + channel.request_pty do |_ch, success| + fail Train::Transports::SSHPTYFailed, 'Requesting PTY failed' unless success + end + end + + channel.exec(cmd) do |_, success| + abort 'Couldn\'t execute command on SSH.' unless success + + channel.on_data do |_, data| + stdout += data + end + + channel.on_extended_data do |_, _type, data| + stderr += data + end + + channel.on_request('exit-status') do |_, data| + exit_status = data.read_long + end + + channel.on_request('exit-signal') do |_, data| + exit_status = data.read_long + end + end + end + @session.loop + + CommandResult.new(stdout, stderr, exit_status) + rescue Net::SSH::Exception => ex + raise Train::Transports::SSHFailed, "SSH command failed (#{ex.message})" + end + # Returns a connection session, or establishes one when invoked the # first time. # diff --git a/lib/train/transports/winrm_connection.rb b/lib/train/transports/winrm_connection.rb index b856d743..291b1d73 100644 --- a/lib/train/transports/winrm_connection.rb +++ b/lib/train/transports/winrm_connection.rb @@ -46,22 +46,6 @@ def close @session = nil end - def file(path) - @files[path] ||= Train::File::Remote::Windows.new(self, path) - end - - def run_command(command) - return if command.nil? - logger.debug("[WinRM] #{self} (#{command})") - out = '' - - response = session.run(command) do |stdout, _| - out << stdout if stdout - end - - CommandResult.new(out, response.stderr, response.exitcode) - end - # (see Base::Connection#login_command) def login_command case RbConfig::CONFIG['host_os'] @@ -107,6 +91,22 @@ def uri PING_COMMAND = "Write-Host '[WinRM] Established\n'".freeze + def file_via_connection(path) + Train::File::Remote::Windows.new(self, path) + end + + def run_command_via_connection(command) + return if command.nil? + logger.debug("[WinRM] #{self} (#{command})") + out = '' + + response = session.run(command) do |stdout, _| + out << stdout if stdout + end + + CommandResult.new(out, response.stderr, response.exitcode) + end + # Create a local RDP document and return it # # @param opts [Hash] configuration options diff --git a/test/unit/plugins/connection_test.rb b/test/unit/plugins/connection_test.rb index a322b55b..21fbef35 100644 --- a/test/unit/plugins/connection_test.rb +++ b/test/unit/plugins/connection_test.rb @@ -10,20 +10,28 @@ connection.close # wont raise end - it 'provides a run_command method' do - proc { connection.run_command('') }.must_raise Train::ClientError + it 'raises an exception for run_command' do + proc { connection.run_command('') }.must_raise NotImplementedError end - it 'provides an os method' do - proc { connection.os }.must_raise Train::ClientError + it 'raises an exception for run_command_via_connection' do + proc { connection.send(:run_command_via_connection, '') }.must_raise NotImplementedError end - it 'provides a file method' do - proc { connection.file('') }.must_raise Train::ClientError + it 'raises an exception for os method' do + proc { connection.os }.must_raise NotImplementedError end - it 'provides a login command method' do - proc { connection.login_command }.must_raise Train::ClientError + it 'raises an exception for file method' do + proc { connection.file('') }.must_raise NotImplementedError + end + + it 'raises an exception for file_via_connection method' do + proc { connection.send(:file_via_connection, '') }.must_raise NotImplementedError + end + + it 'raises an exception for login command method' do + proc { connection.login_command }.must_raise NotImplementedError end it 'can wait until ready' do @@ -40,5 +48,95 @@ cls.new({logger: l}) .method(:logger).call.must_equal(l) end + + describe 'create cache connection' do + it 'default connection cache settings' do + connection.cache_enabled?(:file).must_equal true + connection.cache_enabled?(:command).must_equal false + end + end + + describe 'disable/enable caching' do + it 'disable file cache via connection' do + connection.disable_cache(:file) + connection.cache_enabled?(:file).must_equal false + end + + it 'enable command cache via cache_connection' do + connection.enable_cache(:command) + connection.cache_enabled?(:command).must_equal true + end + + it 'raises an exception for unknown cache type' do + proc { connection.enable_cache(:fake) }.must_raise Train::UnknownCacheType + proc { connection.disable_cache(:fake) }.must_raise Train::UnknownCacheType + end + end + + describe 'cache enable check' do + it 'returns true when cache is enabled' do + cache_enabled = connection.instance_variable_get(:@cache_enabled) + cache_enabled[:test] = true + connection.cache_enabled?(:test).must_equal true + end + + it 'returns false when cache is disabled' do + cache_enabled = connection.instance_variable_get(:@cache_enabled) + cache_enabled[:test] = false + connection.cache_enabled?(:test).must_equal false + end + end + + describe 'clear cache' do + it 'clear file cache' do + cache = connection.instance_variable_get(:@cache) + cache[:file]['/tmp'] = 'test' + connection.send(:clear_cache, :file) + cache = connection.instance_variable_get(:@cache) + cache[:file].must_equal({}) + end + end + + describe 'load file' do + it 'with caching' do + connection.enable_cache(:file) + connection.expects(:file_via_connection).once.returns('test_file') + connection.file('/tmp/test').must_equal('test_file') + connection.file('/tmp/test').must_equal('test_file') + assert = { '/tmp/test' => 'test_file' } + cache = connection.instance_variable_get(:@cache) + cache[:file].must_equal(assert) + end + + it 'without caching' do + connection.disable_cache(:file) + connection.expects(:file_via_connection).twice.returns('test_file') + connection.file('/tmp/test').must_equal('test_file') + connection.file('/tmp/test').must_equal('test_file') + cache = connection.instance_variable_get(:@cache) + cache[:file].must_equal({}) + end + end + + describe 'run command' do + it 'with caching' do + connection.enable_cache(:command) + connection.expects(:run_command_via_connection).once.returns('test_user') + connection.run_command('whoami').must_equal('test_user') + connection.run_command('whoami').must_equal('test_user') + assert = { 'whoami' => 'test_user' } + cache = connection.instance_variable_get(:@cache) + cache[:command].must_equal(assert) + end + + it 'without caching' do + connection.disable_cache(:command) + connection.expects(:run_command_via_connection).twice.returns('test_user') + connection.run_command('whoami').must_equal('test_user') + connection.run_command('whoami').must_equal('test_user') + cache = connection.instance_variable_get(:@cache) + cache[:command].must_equal({}) + end + end end end diff --git a/test/unit/transports/local_test.rb b/test/unit/transports/local_test.rb index 9669c82b..34aae961 100644 --- a/test/unit/transports/local_test.rb +++ b/test/unit/transports/local_test.rb @@ -43,6 +43,16 @@ def initialize connection.login_command.must_be_nil end + it 'provides a run_command_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:run_command_via_connection).must_equal true + end + + it 'provides a file_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:file_via_connection).must_equal true + end + describe 'when running a local command' do let(:cmd_runner) { Minitest::Mock.new } diff --git a/test/unit/transports/mock_test.rb b/test/unit/transports/mock_test.rb index 7b1d4978..ff804589 100644 --- a/test/unit/transports/mock_test.rb +++ b/test/unit/transports/mock_test.rb @@ -19,6 +19,16 @@ connection.uri.must_equal 'mock://' end + it 'provides a run_command_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:run_command_via_connection).must_equal true + end + + it 'provides a file_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:file_via_connection).must_equal true + end + describe 'when running a mocked command' do let(:mock_cmd) { } diff --git a/test/unit/transports/ssh_test.rb b/test/unit/transports/ssh_test.rb index ebb4637f..c9e7a9f5 100644 --- a/test/unit/transports/ssh_test.rb +++ b/test/unit/transports/ssh_test.rb @@ -58,6 +58,16 @@ let(:ssh) { cls.new(conf) } let(:connection) { ssh.connection } + it 'provides a run_command_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:run_command_via_connection).must_equal true + end + + it 'provides a file_via_connection method' do + methods = connection.class.private_instance_methods(false) + methods.include?(:file_via_connection).must_equal true + end + it 'gets the connection' do connection.must_be_kind_of Train::Transports::SSH::Connection end