Skip to content

Improve SSM Session Manager: Add SSH key injection and optimize code #654

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mayankgupta57 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Improve SSM Session Manager: Add SSH key injection and optimize code #654

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3d58819
Improve SSM Session Manager: Add SSH key injection and optimize code
Jan 23, 2026
d161df9
Fix RuboCop offenses: Remove redundant begin, add trailing commas, re…
Jan 23, 2026
4b9bfe1
Add unit tests for SSM Session Manager improvements
Jan 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions docs/ssm-session-manager.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ kitchen-ec2 now supports AWS Systems Manager (SSM) Session Manager as an alterna
- **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
- **Automatic SSH key injection**: SSH keys are automatically injected into instances via SSM before each connection
- **Zero-trust compliance**: Access instances through AWS IAM authentication instead of network-based access

## Requirements
Expand Down Expand Up @@ -61,9 +61,20 @@ suites:
| Option | Default | Description |
|--------|---------|-------------|
| `use_ssm_session_manager` | `false` | Enable SSM Session Manager transport |
| `ssm_session_manager_document_name` | `nil` | Optional custom SSM document name |
| `ssm_session_manager_document_name` | `nil` | Optional custom SSM document name eg. AWS-StartSSHSession |
| `iam_profile_name` | `nil` | IAM instance profile (required for SSM) |

## How It Works

When SSM Session Manager is enabled, kitchen-ec2:

1. Waits for the SSM agent to become available on the instance
2. Automatically injects SSH public keys into the instance via SSM before each connection
3. Configures SSH to use SSM Session Manager as a proxy command
4. Connects to the instance through SSM, eliminating the need for direct network access

SSH keys are automatically extracted from your configured key files and injected into the instance's `~/.ssh/authorized_keys` file using SSM Run Command. This happens transparently before each connection attempt.

## Additional Resources

- [AWS Systems Manager Session Manager Documentation](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager.html)
Expand Down
166 changes: 131 additions & 35 deletions lib/kitchen/driver/ec2.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1282,6 +1282,76 @@ def ssm_session_manager
@ssm_session_manager ||= Aws::SsmSessionManager.new(config, instance.logger)
end

def ssm_session_manager_client
@ssm_session_manager_client ||= ::Aws::SSM::Client.new(
region: config[:region],
profile: config[:shared_credentials_profile]
)
end

def ssm_session_manager_build_proxy_command(state)
# Build SSM start-session command
cmd = [
"aws", "ssm", "start-session",
"--target", state[:server_id],
"--region", config[:region]
]

# Add document name if specified
if config[:ssm_session_manager_document_name]
cmd += ["--document-name", config[:ssm_session_manager_document_name]]
end

# Add AWS profile if specified
if config[:shared_credentials_profile]
cmd += ["--profile", config[:shared_credentials_profile]]
end

cmd.join(" ")
end

def ssm_session_manager_execute_command(state, commands, max_wait: 30)
# Execute SSM command and wait for completion
ssm_client = ssm_session_manager_client

resp = ssm_client.send_command(
instance_ids: [state[:server_id]],
document_name: "AWS-RunShellScript",
parameters: {
"commands" => commands,
}
)

command_id = resp.command.command_id
debug("[AWS SSM Session Manager] Command sent, ID: #{command_id}")

# Wait for command to complete (poll every 2 seconds)
waited = 0
loop do
sleep 2
waited += 2

status_resp = ssm_client.get_command_invocation(
command_id: command_id,
instance_id: state[:server_id]
)

if status_resp.status == "Success"
return true
elsif status_resp.status == "Failed" || status_resp.status == "Cancelled"
warn("[AWS SSM Session Manager] Command failed: #{status_resp.status}")
warn("[AWS SSM Session Manager] Error: #{status_resp.error_message}") if status_resp.error_message
return false
elsif waited >= max_wait
warn("[AWS SSM Session Manager] Timeout waiting for command to complete")
return false
end
end
rescue => e
warn("[AWS SSM Session Manager] Error executing command: #{e.message}")
false
end

def ssm_session_manager_setup_ready(state)
info("[AWS SSM Session Manager] Setting up SSM Session Manager connection")

Expand Down Expand Up @@ -1312,6 +1382,53 @@ def ssm_session_manager_setup_ready(state)
info("[AWS SSM Session Manager] Waiting for SSM agent to be available (attempt #{retries}/#{max_retries})...")
sleep retry_delay
end

# Inject SSH public key into instance via SSM if key exists (one-time setup)
if state[:ssh_key] && File.exist?(state[:ssh_key])
ssm_session_manager_inject_ssh_key(state)
end
end

def ssm_session_manager_inject_ssh_key(state)
# Extract public key using existing helper method
key_path = state[:ssh_key] || instance.transport[:ssh_key]
return unless key_path && File.exist?(key_path)

begin
public_key = instance_connect_extract_public_key(key_path)
rescue => e
warn("[AWS SSM Session Manager] Could not extract public key: #{e.message}")
return
end

return unless public_key

# Use same username resolution pattern as Instance Connect
username = state[:username] || actual_platform&.username
home_dir = username == "root" ? "/root" : "/home/#{username}"
authorized_keys_path = "#{home_dir}/.ssh/authorized_keys"

info("[AWS SSM Session Manager] Injecting SSH public key into instance #{state[:server_id]} for user #{username}")

# Escape the public key for shell
escaped_key = Shellwords.escape(public_key)

# Build SSM command to add key (check if key already exists to avoid duplicates)
# Use test -f to check if file exists before grep to avoid errors
commands = [
"mkdir -p #{home_dir}/.ssh",
"chmod 700 #{home_dir}/.ssh",
"test -f #{authorized_keys_path} && grep -qxF #{escaped_key} #{authorized_keys_path} || echo #{escaped_key} >> #{authorized_keys_path}",
"chmod 600 #{authorized_keys_path}",
"chown #{username}:#{username} #{home_dir}/.ssh -R",
]

# Use reusable SSM command execution helper
if ssm_session_manager_execute_command(state, commands)
info("[AWS SSM Session Manager] SSH public key injected successfully")
else
warn("[AWS SSM Session Manager] Failed to inject SSH key")
end
end

def ssm_session_manager_setup_override(instance)
Expand All @@ -1327,28 +1444,20 @@ def ssm_session_manager_setup_override(instance)

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
# Inject/refresh SSH key before connection (like Instance Connect)
driver_instance.send(:ssm_session_manager_inject_ssh_key, state)

# 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(" ")
# Build and set proxy command
proxy_command = driver_instance.send(:ssm_session_manager_build_proxy_command, state)
driver_instance.info("[AWS SSM Session Manager] Using proxy command: #{proxy_command}")

# Set proxy command for SSH transport
state[:ssh_proxy_command] = proxy_command

# Override hostname to localhost - SSM proxy command handles routing
# SSH should connect to localhost, not the instance IP
driver_instance.debug("[AWS SSM Session Manager] Overriding hostname to localhost (was: #{state[:hostname]})")
state[:hostname] = "localhost"
end

# Call original connection method
Expand Down Expand Up @@ -1388,25 +1497,12 @@ def ssm_session_manager_setup_inspec_override(instance)

# 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
# Inject/refresh SSH key before connection (like Instance Connect)
driver_instance.send(:ssm_session_manager_inject_ssh_key, config_data)

opts["proxy_command"] = cmd.join(" ")
driver_instance.info("[AWS SSM Session Manager] InSpec using proxy command: #{opts["proxy_command"]}")
proxy_command = driver_instance.send(:ssm_session_manager_build_proxy_command, config_data)
opts["proxy_command"] = proxy_command
driver_instance.info("[AWS SSM Session Manager] InSpec using proxy command: #{proxy_command}")
end

opts
Expand Down
Loading