diff --git a/.gitignore b/.gitignore index c5807ea0..cf52a30b 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,4 @@ tmp .kitchen/ .kitchen.local.yml +vendor/bundle/ diff --git a/README.md b/README.md index 72b5752f..eeb925cb 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,16 @@ will need access to an [AWS][aws_site] account. [IAM][iam_site] users should hav By automatically applying reasonable defaults wherever possible, kitchen-ec2 does a lot of work to make your life easier. See the [kitchen.ci kitchen-ec2 docs](https://kitchen.ci/docs/drivers/aws/) for a complete list of configuration options. +### Transport Options + +kitchen-ec2 supports multiple transport methods for connecting to EC2 instances: + +- **SSH/WinRM** (default): Traditional network-based access +- **Instance Connect**: AWS EC2 Instance Connect for temporary SSH access +- **SSM Session Manager**: AWS Systems Manager Session Manager for secure, audited access without SSH/RDP ports + +For SSM Session Manager configuration and benefits, see [SSM Session Manager Documentation](docs/ssm-session-manager.md). + ## Development * Source hosted at [GitHub][repo] diff --git a/docs/ssm-session-manager.md b/docs/ssm-session-manager.md new file mode 100644 index 00000000..158e1cac --- /dev/null +++ b/docs/ssm-session-manager.md @@ -0,0 +1,70 @@ +# AWS SSM Session Manager Support + +kitchen-ec2 now supports AWS Systems Manager (SSM) Session Manager as an alternative transport method to SSH/WinRM. This feature enables Test Kitchen to connect to EC2 instances without requiring direct network connectivity or SSH key management. + +## Benefits + +- **No SSH/WinRM network access required**: Connect to instances in private subnets without VPN or bastion hosts +- **Enhanced security**: No need to open SSH/RDP ports in security groups +- **Centralized audit logging**: All session activity is logged to CloudTrail +- **No SSH key management**: Eliminate the complexity of managing SSH key pairs for testing +- **Zero-trust compliance**: Access instances through AWS IAM authentication instead of network-based access + +## Requirements + +### Client Requirements + +1. **AWS CLI**: Install the AWS CLI version 2.x or later +2. **Session Manager Plugin**: Install the Session Manager plugin for AWS CLI + +### Instance Requirements + +1. **SSM Agent**: Must be installed and running on the EC2 instance +2. **IAM Instance Profile**: Instance must have the `AmazonSSMManagedInstanceCore` managed policy or equivalent +3. **Network Access**: Outbound HTTPS (port 443) access to AWS SSM endpoints + +## Configuration + +### Basic Configuration + +```yaml +driver: + name: ec2 + use_ssm_session_manager: true + iam_profile_name: my-ssm-enabled-profile +``` + +### Complete Example + +```yaml +driver: + name: ec2 + use_ssm_session_manager: true + instance_type: t3.micro + subnet_id: subnet-12345678 + iam_profile_name: kitchen-ec2-ssm-profile + security_group_ids: + - sg-87654321 + +platforms: + - name: amazon2 + - name: ubuntu-20.04 + +suites: + - name: default + run_list: + - recipe[my_cookbook::default] +``` + +## Configuration Options + +| Option | Default | Description | +|--------|---------|-------------| +| `use_ssm_session_manager` | `false` | Enable SSM Session Manager transport | +| `ssm_session_manager_document_name` | `nil` | Optional custom SSM document name | +| `iam_profile_name` | `nil` | IAM instance profile (required for SSM) | + +## Additional Resources + +- [AWS Systems Manager Session Manager Documentation](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager.html) +- [Session Manager Plugin Installation](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html) diff --git a/kitchen-ec2.gemspec b/kitchen-ec2.gemspec index caed1356..df27c4de 100644 --- a/kitchen-ec2.gemspec +++ b/kitchen-ec2.gemspec @@ -19,6 +19,7 @@ Gem::Specification.new do |gem| gem.add_dependency "aws-sdk-ec2", "~> 1.0" gem.add_dependency "aws-sdk-ec2instanceconnect", "~> 1.0" + gem.add_dependency "aws-sdk-ssm", "~> 1.0" gem.add_dependency "retryable", ">= 2.0", "< 4.0" # 4.0 will need to be validated gem.add_dependency "sshkey", "~> 2.0" gem.add_dependency "test-kitchen", ">= 3.9.0", "< 4" diff --git a/lib/kitchen/driver/aws/ssm_session_manager.rb b/lib/kitchen/driver/aws/ssm_session_manager.rb new file mode 100644 index 00000000..b6a629e1 --- /dev/null +++ b/lib/kitchen/driver/aws/ssm_session_manager.rb @@ -0,0 +1,85 @@ +# +# Author:: GitHub Copilot +# +# Copyright:: 2025, GitHub +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +require "aws-sdk-ssm" +require "open3" + +module Kitchen + module Driver + class Aws + # Manages AWS Systems Manager Session Manager connections for Test Kitchen + class SsmSessionManager + def initialize(config, logger) + @config = config + @logger = logger + @ssm_client = ::Aws::SSM::Client.new( + region: config[:region], + profile: config[:shared_credentials_profile] + ) + end + + # Check if SSM agent is running on the instance + def ssm_agent_available?(instance_id) + @logger.debug("Checking if SSM agent is available on instance #{instance_id}") + + begin + resp = @ssm_client.describe_instance_information( + filters: [ + { + key: "InstanceIds", + values: [instance_id], + }, + ] + ) + + available = !resp.instance_information_list.empty? && + resp.instance_information_list.first.ping_status == "Online" + + if available + @logger.info("SSM agent is available on instance #{instance_id}") + else + @logger.warn("SSM agent is not available on instance #{instance_id}") + end + + available + rescue ::Aws::SSM::Errors::ServiceError => e + @logger.warn("Error checking SSM agent status: #{e.message}") + false + end + end + + # Verify that the AWS CLI session manager plugin is installed + def session_manager_plugin_installed? + _output, status = Open3.capture2e("session-manager-plugin", "--version") + installed = status.success? + + if installed + @logger.debug("Session Manager plugin is installed") + else + @logger.warn("Session Manager plugin is not installed. Install it from: " \ + "https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html") + end + + installed + rescue StandardError => e + @logger.warn("Error checking for session-manager-plugin: #{e.message}") + false + end + end + end + end +end diff --git a/lib/kitchen/driver/ec2.rb b/lib/kitchen/driver/ec2.rb index 628c0598..2947beee 100644 --- a/lib/kitchen/driver/ec2.rb +++ b/lib/kitchen/driver/ec2.rb @@ -39,6 +39,7 @@ require_relative "aws/standard_platform/ubuntu" require_relative "aws/standard_platform/windows" require_relative "aws/instance_connect" +require_relative "aws/ssm_session_manager" require "aws-sdk-ec2" require "aws-sdk-core/waiters/errors" require "retryable" unless defined?(Retryable) @@ -101,6 +102,8 @@ class Ec2 < Kitchen::Driver::Base default_config :use_instance_connect, false default_config :instance_connect_endpoint_id, nil default_config :instance_connect_max_tunnel_duration, 3600 + default_config :use_ssm_session_manager, false + default_config :ssm_session_manager_document_name, nil include Kitchen::Driver::Mixins::DedicatedHosts @@ -197,6 +200,15 @@ def self.validation_error(driver, old_key, new_key) end end + # Ensure use_instance_connect and use_ssm_session_manager are not both enabled + validations[:use_ssm_session_manager] = lambda do |_attr, val, driver| + if val && driver[:use_instance_connect] + warn "Cannot use both 'use_instance_connect' and 'use_ssm_session_manager' at the same time. " \ + "Please enable only one transport method." + exit! + end + end + # empty keys cause failures when tagging and they make no sense validations[:tags] = lambda do |_attr, val, _driver| # if someone puts the tags each on their own line it's an array not a hash @@ -272,7 +284,7 @@ def create(state) # Waiting can fail, so we have to retry on that. Retryable.retryable( tries: 10, - sleep: lambda { |n| [2**n, 30].min }, + sleep: ->(n) { [2**n, 30].min }, on: ::Aws::EC2::Errors::InvalidInstanceIDNotFound ) do |_r, _| wait_until_ready(server, state) @@ -282,6 +294,8 @@ def create(state) if config[:use_instance_connect] instance_connect_setup_ready(state) + elsif config[:use_ssm_session_manager] + ssm_session_manager_setup_ready(state) end instance.transport.connection(state).wait_until_ready @@ -456,10 +470,7 @@ def submit_spots expanded = [] keys = %i{instance_type} - unless config[:subnet_filter] - # => Use explicitly specified subnets - keys << :subnet_id - else + if config[:subnet_filter] # => Enable cascading through matching subnets client = ::Aws::EC2::Client.new(region: config[:region]) @@ -484,6 +495,9 @@ def submit_spots new_config[:subnet_filter] = nil new_config end + else + # => Use explicitly specified subnets + keys << :subnet_id end keys.each do |key| @@ -538,7 +552,7 @@ def submit_spot # not retry if the price could not be satisfied immediately. Retryable.retryable( tries: config[:spot_wait] / config[:retryable_sleep], - sleep: lambda { |_n| config[:retryable_sleep] }, + sleep: ->(_n) { config[:retryable_sleep] }, on: ::Aws::EC2::Errors::SpotMaxPriceTooLow ) do |retries| c = retries * config[:retryable_sleep] @@ -577,7 +591,7 @@ def wait_until_ready(server, state) output = Base64.decode64(output) debug "Console output: --- \n#{output}" end - ready = !!(output.include?("Windows is Ready to use")) + ready = !!output.include?("Windows is Ready to use") end end ready @@ -683,7 +697,7 @@ def sudo_command def create_ec2_json(state) if windows_os? - cmd = "New-Item -Force C:\\chef\\ohai\\hints\\ec2.json -ItemType File" + cmd = 'New-Item -Force C:\\chef\\ohai\\hints\\ec2.json -ItemType File' else debug "Using sudo_command='#{sudo_command}' for ohai hints" cmd = "#{sudo_command} mkdir -p /etc/chef/ohai/hints; #{sudo_command} touch /etc/chef/ohai/hints/ec2.json" @@ -911,16 +925,16 @@ def attach_network_interface(state) client = ::Aws::EC2::Client.new(region: config[:region]) begin check_eni = client.describe_network_interface_attribute({ - attribute: "attachment", - network_interface_id: config[:elastic_network_interface_id], - }) + attribute: "attachment", + network_interface_id: config[:elastic_network_interface_id], + }) if check_eni.attachment.nil? unless state[:server_id].nil? client.attach_network_interface({ - device_index: 1, - instance_id: state[:server_id], - network_interface_id: config[:elastic_network_interface_id], - }) + device_index: 1, + instance_id: state[:server_id], + network_interface_id: config[:elastic_network_interface_id], + }) info("Attached Network interface <#{config[:elastic_network_interface_id]}> with the instance <#{state[:server_id]}> .") end else @@ -966,6 +980,10 @@ def finalize_config!(instance) debug("[AWS EC2 Instance Connect] Setting up Instance Connect overrides") instance_connect_setup_override(instance) instance_connect_setup_inspec_override(instance) + elsif config[:use_ssm_session_manager] + debug("[AWS SSM Session Manager] Setting up SSM Session Manager overrides") + ssm_session_manager_setup_override(instance) + ssm_session_manager_setup_inspec_override(instance) end self @@ -1037,7 +1055,6 @@ def instance_connect_setup_inspec_override(instance) # Override runner_options_for_ssh define_singleton_method(:runner_options_for_ssh) do |config_data| - # Get the original options opts = instance_connect_original_runner_options_for_ssh(config_data) @@ -1076,8 +1093,8 @@ def instance_connect_setup_inspec_override(instance) if public_dns && !public_dns.empty? opts["host"] = public_dns opts["ssh_options"] = (opts["ssh_options"] || {}).merge({ - "IdentitiesOnly" => "yes", - }) + "IdentitiesOnly" => "yes", + }) driver_instance.info("[AWS EC2 Instance Connect] InSpec using direct SSH to #{public_dns} with IdentitiesOnly=yes") else driver_instance.warn("[AWS EC2 Instance Connect] No public DNS available for direct SSH mode") @@ -1207,11 +1224,11 @@ def instance_connect_endpoint_available?(state) def get_vpc_id_for_instance(state) # Get the instance details to find its VPC - return nil unless state[:server_id] + return unless state[:server_id] begin instance_info = ec2.client.describe_instances(instance_ids: [state[:server_id]]).reservations.first&.instances&.first - return nil unless instance_info + return unless instance_info instance_info.vpc_id rescue => e @@ -1258,6 +1275,152 @@ def instance_connect_extract_public_key(private_key_path) raise "Unable to extract public key from #{private_key_path}: #{e.message}" end end + + # SSM Session Manager Support Methods + + def ssm_session_manager + @ssm_session_manager ||= Aws::SsmSessionManager.new(config, instance.logger) + end + + def ssm_session_manager_setup_ready(state) + info("[AWS SSM Session Manager] Setting up SSM Session Manager connection") + + # Verify session manager plugin is installed + unless ssm_session_manager.session_manager_plugin_installed? + warn("[AWS SSM Session Manager] Session Manager plugin not found. Please install it from: " \ + "https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html") + end + + # Wait for SSM agent to be available + max_retries = 12 + retry_delay = 10 + retries = 0 + + loop do + if ssm_session_manager.ssm_agent_available?(state[:server_id]) + info("[AWS SSM Session Manager] SSM agent is available on instance #{state[:server_id]}") + break + end + + retries += 1 + if retries >= max_retries + warn("[AWS SSM Session Manager] SSM agent did not become available after #{max_retries * retry_delay} seconds. " \ + "Ensure the instance has an IAM instance profile with SSM permissions and the SSM agent is running.") + break + end + + info("[AWS SSM Session Manager] Waiting for SSM agent to be available (attempt #{retries}/#{max_retries})...") + sleep retry_delay + end + end + + def ssm_session_manager_setup_override(instance) + # Prevent double setup + return if instance.transport.respond_to?(:ssm_session_manager_override_applied) + + # Store reference to driver for use in override + driver_instance = self + use_ssm = config[:use_ssm_session_manager] + + # Override the transport's connection method to inject SSM setup + original_connection = instance.transport.method(:connection) + + instance.transport.define_singleton_method(:connection) do |state, &block| + if use_ssm && state[:server_id] + # Build SSM start-session command + cmd = [ + "aws", "ssm", "start-session", + "--target", state[:server_id], + "--region", driver_instance.config[:region] + ] + + # Add document name if specified + if driver_instance.config[:ssm_session_manager_document_name] + cmd += ["--document-name", driver_instance.config[:ssm_session_manager_document_name]] + end + + # Add AWS profile if specified + if driver_instance.config[:shared_credentials_profile] + cmd += ["--profile", driver_instance.config[:shared_credentials_profile]] + end + + proxy_command = cmd.join(" ") + driver_instance.info("[AWS SSM Session Manager] Using proxy command: #{proxy_command}") + + # Set proxy command for SSH transport + state[:ssh_proxy_command] = proxy_command + end + + # Call original connection method + original_connection.call(state, &block) + end + + # Mark as applied + instance.transport.define_singleton_method(:ssm_session_manager_override_applied) { true } + end + + def ssm_session_manager_setup_inspec_override(instance) + # Only apply to InSpec verifier + return unless instance.verifier.name.downcase == "inspec" + return if instance.verifier.respond_to?(:ssm_session_manager_inspec_override_applied) + + # Store reference to driver for use in override + driver_instance = self + use_ssm = config[:use_ssm_session_manager] + + # Override the verifier's call method to inject SSM setup + original_call = instance.verifier.method(:call) + + instance.verifier.define_singleton_method(:call) do |state| + driver_instance.debug("[AWS SSM Session Manager] InSpec call method intercepted") + + # Set up SSM for InSpec if enabled + if use_ssm && state[:server_id] + # Check if we already have the override method defined + unless respond_to?(:ssm_original_runner_options_for_ssh) + # Store the original method + define_singleton_method(:ssm_original_runner_options_for_ssh, method(:runner_options_for_ssh)) + + # Override runner_options_for_ssh + define_singleton_method(:runner_options_for_ssh) do |config_data| + # Get the original options + opts = ssm_original_runner_options_for_ssh(config_data) + + # Inject SSM Session Manager configuration if enabled + if use_ssm && config_data[:server_id] + # Build SSM start-session command + cmd = [ + "aws", "ssm", "start-session", + "--target", config_data[:server_id], + "--region", driver_instance.config[:region] + ] + + # Add document name if specified + if driver_instance.config[:ssm_session_manager_document_name] + cmd += ["--document-name", driver_instance.config[:ssm_session_manager_document_name]] + end + + # Add AWS profile if specified + if driver_instance.config[:shared_credentials_profile] + cmd += ["--profile", driver_instance.config[:shared_credentials_profile]] + end + + opts["proxy_command"] = cmd.join(" ") + driver_instance.info("[AWS SSM Session Manager] InSpec using proxy command: #{opts["proxy_command"]}") + end + + opts + end + end + end + + # Call the original method + original_call.call(state) + end + + # Mark as applied + instance.verifier.define_singleton_method(:ssm_session_manager_inspec_override_applied) { true } + end end end end diff --git a/spec/kitchen/driver/aws/image_selection_spec.rb b/spec/kitchen/driver/aws/image_selection_spec.rb index 0a1c53e7..9825ac98 100644 --- a/spec/kitchen/driver/aws/image_selection_spec.rb +++ b/spec/kitchen/driver/aws/image_selection_spec.rb @@ -160,7 +160,7 @@ def new_instance(platform_name: "blarghle") "debian" => [ { name: "owner-id", values: %w{136693071363} }, - { name: "name", values: %w{debian-11-*} }, + { name: "name", values: %w{debian-13-*} }, ], "debian-13" => [ { name: "owner-id", values: %w{136693071363} }, @@ -188,7 +188,7 @@ def new_instance(platform_name: "blarghle") ], "debian-x86_64" => [ { name: "owner-id", values: %w{136693071363} }, - { name: "name", values: %w{debian-11-*} }, + { name: "name", values: %w{debian-13-*} }, { name: "architecture", values: %w{x86_64} }, ], diff --git a/spec/kitchen/driver/aws/ssm_session_manager_spec.rb b/spec/kitchen/driver/aws/ssm_session_manager_spec.rb new file mode 100644 index 00000000..431d0a01 --- /dev/null +++ b/spec/kitchen/driver/aws/ssm_session_manager_spec.rb @@ -0,0 +1,110 @@ +# +# Author:: GitHub Copilot +# +# Copyright:: 2025, GitHub +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +require 'kitchen/driver/aws/ssm_session_manager' + +describe Kitchen::Driver::Aws::SsmSessionManager do + let(:config) do + { + region: 'us-west-2', + shared_credentials_profile: nil, + } + end + let(:logger) { instance_double(Logger, debug: nil, info: nil, warn: nil) } + let(:ssm_client) { instance_double(Aws::SSM::Client) } + let(:ssm_manager) { described_class.new(config, logger) } + + before do + allow(Aws::SSM::Client).to receive(:new).and_return(ssm_client) + end + + describe '#ssm_agent_available?' do + let(:instance_id) { 'i-1234567890abcdef0' } + + context 'when SSM agent is online' do + it 'returns true' do + response = double( + instance_information_list: [ + double(ping_status: 'Online'), + ] + ) + expect(ssm_client).to receive(:describe_instance_information) + .with(filters: [{ key: 'InstanceIds', values: [instance_id] }]) + .and_return(response) + + expect(ssm_manager.ssm_agent_available?(instance_id)).to be true + end + end + + context 'when SSM agent is not online' do + it 'returns false' do + response = double( + instance_information_list: [ + double(ping_status: 'ConnectionLost'), + ] + ) + expect(ssm_client).to receive(:describe_instance_information) + .with(filters: [{ key: 'InstanceIds', values: [instance_id] }]) + .and_return(response) + + expect(ssm_manager.ssm_agent_available?(instance_id)).to be false + end + end + + context 'when instance not found in SSM' do + it 'returns false' do + response = double(instance_information_list: []) + expect(ssm_client).to receive(:describe_instance_information) + .with(filters: [{ key: 'InstanceIds', values: [instance_id] }]) + .and_return(response) + + expect(ssm_manager.ssm_agent_available?(instance_id)).to be false + end + end + + context 'when SSM API call fails' do + it 'returns false and logs warning' do + expect(ssm_client).to receive(:describe_instance_information) + .and_raise(Aws::SSM::Errors::ServiceError.new(nil, 'API Error')) + + expect(logger).to receive(:warn).with(/Error checking SSM agent status/) + expect(ssm_manager.ssm_agent_available?(instance_id)).to be false + end + end + end + + describe '#session_manager_plugin_installed?' do + context 'when plugin is installed' do + it 'returns true' do + status_double = double(success?: true) + allow(Open3).to receive(:capture2e).with('session-manager-plugin', '--version').and_return(["1.2.3\n", status_double]) + + expect(ssm_manager.session_manager_plugin_installed?).to be true + end + end + + context 'when plugin is not installed' do + it 'returns false and logs warning' do + status_double = double(success?: false) + allow(Open3).to receive(:capture2e).with('session-manager-plugin', '--version').and_return(['', status_double]) + + expect(logger).to receive(:warn).with(/Session Manager plugin is not installed/) + expect(ssm_manager.session_manager_plugin_installed?).to be false + end + end + end +end diff --git a/spec/kitchen/driver/ec2_spec.rb b/spec/kitchen/driver/ec2_spec.rb index fbd4cc1c..60c258ad 100644 --- a/spec/kitchen/driver/ec2_spec.rb +++ b/spec/kitchen/driver/ec2_spec.rb @@ -1268,4 +1268,155 @@ end end end + + describe "SSM Session Manager functionality" do + describe "integration with create method" do + let(:ssm_manager) { instance_double(Kitchen::Driver::Aws::SsmSessionManager) } + let(:create_state) { {} } + let(:config) do + { + use_ssm_session_manager: true, + region: "us-east-1", + aws_ssh_key_id: "test-key", + image_id: "ami-12345", + security_group_ids: ["sg-12345"], + } + end + + before do + allow(driver).to receive(:update_username) + allow(driver).to receive(:submit_server).and_return(server) + allow(server).to receive(:wait_until_exists) + allow(server).to receive(:id).and_return("i-12345") + allow(driver).to receive(:attach_network_interface) + allow(driver).to receive(:create_ec2_json) + allow(driver).to receive(:debug) + allow(driver).to receive(:info) + allow(driver).to receive(:create_security_group) + allow(driver).to receive(:create_key) + allow(driver).to receive(:host_available?).and_return(true) + allow(driver).to receive(:ssm_session_manager).and_return(ssm_manager) + + allow(driver).to receive(:wait_until_ready) do |_, state| + state[:hostname] = "test-hostname" + true + end + + allow(Retryable).to receive(:retryable) do |_options, &block| + block.call(1, nil) + end + + connection_double = double("connection", wait_until_ready: true) + allow(instance.transport).to receive(:connection).with(any_args).and_return(connection_double) + end + + it "calls ssm_session_manager_setup_ready when use_ssm_session_manager is true" do + expect(driver).to receive(:ssm_session_manager_setup_ready).with(hash_including(server_id: "i-12345")) + driver.create(create_state) + end + + it "does not call ssm_session_manager_setup_ready when use_ssm_session_manager is false" do + config[:use_ssm_session_manager] = false + expect(driver).not_to receive(:ssm_session_manager_setup_ready) + driver.create(create_state) + end + end + + describe "#ssm_session_manager_setup_ready" do + let(:state) { { server_id: "i-12345" } } + let(:ssm_manager) { instance_double(Kitchen::Driver::Aws::SsmSessionManager) } + + before do + allow(driver).to receive(:ssm_session_manager).and_return(ssm_manager) + allow(driver).to receive(:info) + allow(driver).to receive(:warn) + allow(ssm_manager).to receive(:session_manager_plugin_installed?).and_return(true) + end + + context "when SSM agent is available immediately" do + it "does not wait" do + expect(ssm_manager).to receive(:ssm_agent_available?).with("i-12345").and_return(true) + expect(driver).not_to receive(:sleep) + driver.send(:ssm_session_manager_setup_ready, state) + end + end + + context "when SSM agent becomes available after retries" do + it "waits and succeeds" do + call_count = 0 + allow(ssm_manager).to receive(:ssm_agent_available?) do |_instance_id| + call_count += 1 + call_count > 2 + end + + expect(driver).to receive(:sleep).twice + driver.send(:ssm_session_manager_setup_ready, state) + end + end + + context "when session manager plugin is not installed" do + it "warns the user" do + allow(ssm_manager).to receive(:session_manager_plugin_installed?).and_return(false) + allow(ssm_manager).to receive(:ssm_agent_available?).and_return(true) + + expect(driver).to receive(:warn).with(/Session Manager plugin not found/) + driver.send(:ssm_session_manager_setup_ready, state) + end + end + end + + describe "#ssm_session_manager_setup_override" do + let(:state) { { server_id: "i-12345" } } + let(:config) do + { + use_ssm_session_manager: true, + region: "us-west-2", + shared_credentials_profile: "my-profile", + } + end + + before do + allow(driver).to receive(:info) + end + + it "overrides transport connection method" do + expect(instance.transport).to receive(:define_singleton_method).with(:connection) + expect(instance.transport).to receive(:define_singleton_method).with(:ssm_session_manager_override_applied) + driver.send(:ssm_session_manager_setup_override, instance) + end + + it "sets SSH proxy command with correct parameters" do + driver.send(:ssm_session_manager_setup_override, instance) + + # Call the overridden connection method + original_method = instance.transport.method(:connection) + allow(original_method).to receive(:call) + + instance.transport.connection(state) + + expect(state[:ssh_proxy_command]).to include("aws ssm start-session") + expect(state[:ssh_proxy_command]).to include("--target i-12345") + expect(state[:ssh_proxy_command]).to include("--region us-west-2") + expect(state[:ssh_proxy_command]).to include("--profile my-profile") + end + + it "includes document name when specified" do + config[:ssm_session_manager_document_name] = "MyCustomDocument" + driver.send(:ssm_session_manager_setup_override, instance) + + instance.transport.connection(state) + + expect(state[:ssh_proxy_command]).to include("--document-name MyCustomDocument") + end + + it "does not apply override twice" do + driver.send(:ssm_session_manager_setup_override, instance) + expect(instance.transport).to respond_to(:ssm_session_manager_override_applied) + + # Second call should return early + expect(instance.transport).not_to receive(:define_singleton_method) + driver.send(:ssm_session_manager_setup_override, instance) + end + end + end end