|
| 1 | +require "aws-sdk-ec2" |
| 2 | +require "aws-sdk-ssm" |
| 3 | +require "resolv" |
| 4 | +require "train" |
| 5 | + |
| 6 | +module TrainPlugins |
| 7 | + module AWSSSM |
| 8 | + class Connection < Train::Plugins::Transport::BaseConnection |
| 9 | + def initialize(options) |
| 10 | + super(options) |
| 11 | + |
| 12 | + @ssm = Aws::SSM::Client.new |
| 13 | + end |
| 14 | + |
| 15 | + def close |
| 16 | + logger.info format("[AWS-SSM] Closed connection to %s", @options[:host]) |
| 17 | + end |
| 18 | + |
| 19 | + def uri |
| 20 | + "aws-ssm://#{@options[:host]}/" |
| 21 | + end |
| 22 | + |
| 23 | + def run_command_via_connection(cmd, &data_handler) |
| 24 | + logger.info format("[AWS-SSM] Sending command to %s", @options[:host]) |
| 25 | + exit_status, stdout, stderr = execute_on_channel(cmd, &data_handler) |
| 26 | + |
| 27 | + CommandResult.new(stdout, stderr, exit_status) |
| 28 | + end |
| 29 | + |
| 30 | + def execute_on_channel(cmd, &data_handler) |
| 31 | + logger.debug format("[AWS-SSM] Command: '%s'", cmd) |
| 32 | + |
| 33 | + result = execute_command(@options[:host], cmd) |
| 34 | + |
| 35 | + stdout = result.standard_output_content || "" |
| 36 | + stderr = result.standard_error_content || "" |
| 37 | + exit_status = result.response_code |
| 38 | + |
| 39 | + [exit_status, stdout, stderr] |
| 40 | + end |
| 41 | + |
| 42 | + private |
| 43 | + |
| 44 | + # Check if this is an IP address |
| 45 | + def ip_address?(address) |
| 46 | + !!(address =~ Resolv::IPv4::Regex) |
| 47 | + end |
| 48 | + |
| 49 | + # Check if this is a DNS name |
| 50 | + def dns_name?(address) |
| 51 | + !ip_address?(address) |
| 52 | + end |
| 53 | + |
| 54 | + # Check if this is an internal/external AWS DNS entry |
| 55 | + def amazon_dns?(dns) |
| 56 | + dns.end_with?(".compute.amazonaws.com") || dns.end_with?(".compute.internal") |
| 57 | + end |
| 58 | + |
| 59 | + # Resolve EC2 instance ID associated with a primary IP or a DNS entry |
| 60 | + def instance_id(address) |
| 61 | + logger.debug format("[AWS-SSM] Trying to resolve address %s", address) |
| 62 | + |
| 63 | + ec2 = Aws::EC2::Client.new |
| 64 | + instances = ec2.describe_instances.reservations.collect { |r| r.instances.first } |
| 65 | + |
| 66 | + # Resolve, if DNS name and not Amazon default |
| 67 | + if dns_name?(address) && !amazon_dns?(address) |
| 68 | + address = Resolv.getaddress(address) |
| 69 | + logger.debug format("[AWS-SSM] Resolved non-internal AWS address to %s", address) |
| 70 | + end |
| 71 | + |
| 72 | + # Check the primary IPs and hostnames for a match |
| 73 | + id = instances.detect do |i| |
| 74 | + [ |
| 75 | + i.private_ip_address, |
| 76 | + i.public_ip_address, |
| 77 | + i.private_dns_name, |
| 78 | + i.public_dns_name, |
| 79 | + ].include?(address) |
| 80 | + end&.instance_id |
| 81 | + |
| 82 | + raise format("Could not resolve instance ID for address %s", address) if id.nil? |
| 83 | + |
| 84 | + logger.debug format("[AWS-SSM] Resolved address %s to instance ID %s", address, id) |
| 85 | + id |
| 86 | + end |
| 87 | + |
| 88 | + # Request a command invocation and wait until it is registered with an ID |
| 89 | + def wait_for_invocation(instance_id, command_id) |
| 90 | + invocation_result(instance_id, command_id) |
| 91 | + |
| 92 | + # Retry until the invocation was created on AWS |
| 93 | + rescue Aws::SSM::Errors::InvocationDoesNotExist |
| 94 | + sleep @options[:recheck_invocation] |
| 95 | + retry |
| 96 | + end |
| 97 | + |
| 98 | + # Return the result of a given command invocation |
| 99 | + def invocation_result(instance_id, command_id) |
| 100 | + @ssm.get_command_invocation(instance_id: instance_id, command_id: command_id) |
| 101 | + end |
| 102 | + |
| 103 | + # Return if a non-terminal command status was given |
| 104 | + # @see https://docs.aws.amazon.com/systems-manager/latest/userguide/monitor-commands.html |
| 105 | + def in_progress?(name) |
| 106 | + %w{Pending InProgress Delayed}.include? name |
| 107 | + end |
| 108 | + |
| 109 | + # Return if a terminal command status was given |
| 110 | + # @see https://docs.aws.amazon.com/systems-manager/latest/userguide/monitor-commands.html |
| 111 | + def terminal_state?(name) |
| 112 | + !in_progress?(name) |
| 113 | + end |
| 114 | + |
| 115 | + # Execute a command via SSM |
| 116 | + def execute_command(address, command) |
| 117 | + instance_id = instance_id(address) |
| 118 | + |
| 119 | + cmd = @ssm.send_command(instance_ids: [instance_id], document_name: "AWS-RunShellScript", parameters: { "commands": [command] }) |
| 120 | + cmd_id = cmd.command.command_id |
| 121 | + |
| 122 | + wait_for_invocation(instance_id, cmd_id) |
| 123 | + logger.debug format("[AWS-SSM] Execution ID %s", cmd_id) |
| 124 | + |
| 125 | + start_time = Time.now |
| 126 | + result = invocation_result(instance_id, cmd.command.command_id) |
| 127 | + |
| 128 | + until terminal_state?(result.status) || Time.now - start_time > @options[:execution_timeout] |
| 129 | + result = invocation_result(instance_id, cmd.command.command_id) |
| 130 | + sleep @options[:recheck_execution] |
| 131 | + end |
| 132 | + |
| 133 | + if Time.now - start_time > @options[:execution_timeout] |
| 134 | + raise format("Timeout waiting for execution") |
| 135 | + elsif result.status != "Success" |
| 136 | + raise format('Execution failed with state "%s": %s', result.status, result.standard_error_content || "unknown") |
| 137 | + end |
| 138 | + |
| 139 | + result |
| 140 | + end |
| 141 | + end |
| 142 | + end |
| 143 | +end |
0 commit comments