Generate temporary security groups in AWS

README.md

@@ -45,7 +45,8 @@ Options:
   -i, --aws-instance-type=<s>           The ec2 instance type, defaults to t2.micro (default:
                                         t2.micro)
   -c, --aws-security-groups=<s>         Security groups to associate to the launched instances. May be
-                                        specified multiple times
+                                        specified multiple times. If not provided a temporary security
+                                        group will be generated in AWS
   -p, --aws-public-ip                   Launch instances with a public IP
   -t, --ssh-retries=<i>                 The number of times we should try sshing to the ec2 instance
                                         before giving up. Defaults to 30 (default: 30)

lib/ami_spec.rb

@@ -1,6 +1,7 @@
 require 'ami_spec/aws_instance'
 require 'ami_spec/aws_instance_options'
 require 'ami_spec/aws_key_pair'
+require 'ami_spec/aws_security_group'
 require 'ami_spec/server_spec'
 require 'ami_spec/server_spec_options'
 require 'ami_spec/wait_for_ssh'
@@ -29,7 +30,9 @@ class InstanceConnectionTimeout < StandardError; end
   #   Defaults to AWS_DEFAULT_REGION
   # aws_security_group_ids::
   #   AWS Security groups to assign to the instances
-  #   Defaults to the default security group for the VPC
+  #   If not provided a temporary security group will be generated in AWS
+  # allow_any_temporary_security_group::
+  #   The temporary security group will allow SSH connections from any IP address (0.0.0.0/0)
   # aws_instance_type::
   #   AWS ec2 instance type
   # aws_public_ip::
@@ -57,6 +60,16 @@ def self.run(options)
       options[:key_file] = key_pair.key_file
     end
 
+    if options[:aws_security_groups].nil? || options[:aws_security_groups].empty?
+      temporary_security_group = AwsSecurityGroup.create(
+        ec2: ec2,
+        subnet_id: options[:subnet_id],
+        allow_any_ip: options[:allow_any_temporary_security_group],
+        logger: logger
+      )
+      options[:aws_security_groups] = [temporary_security_group.group_id]
+    end
+
     instances = []
     options[:amis].each_pair do |role, ami|
       aws_instance_options = AwsInstanceOptions.new(options.merge(role: role, ami: ami))
@@ -77,6 +90,7 @@ def self.run(options)
   ensure
     stop_instances(instances, options[:debug])
     key_pair.delete if key_pair
+    temporary_security_group.delete if temporary_security_group
   end
 
   def self.stop_instances(instances, debug)
@@ -109,8 +123,9 @@ def self.invoke
       opt :ssh_user, "The user to ssh to the instance as", type: :string, required: true
       opt :aws_region, "The AWS region, defaults to AWS_DEFAULT_REGION environment variable", type: :string
       opt :aws_instance_type, "The ec2 instance type, defaults to t2.micro", type: :string, default: 't2.micro'
-      opt :aws_security_groups, "Security groups to associate to the launched instances. May be specified multiple times",
+      opt :aws_security_groups, "Security groups to associate to the launched instances. May be specified multiple times. If not provided a temporary security group will be generated in AWS",
           type: :string, default: nil, multi: true
+      opt :allow_any_temporary_security_group, "The temporary security group will allow SSH connections from any IP address (0.0.0.0/0)"
       opt :aws_public_ip, "Launch instances with a public IP"
       opt :ssh_retries, "The number of times we should try sshing to the ec2 instance before giving up. Defaults to 30",
           type: :int, default: 30

lib/ami_spec/aws_instance.rb

@@ -26,7 +26,7 @@ def initialize(options)
       @iam_instance_profile_arn = options.fetch(:iam_instance_profile_arn, nil)
     end
 
-    def_delegators :@instance, :instance_id, :tags, :terminate, :private_ip_address, :public_ip_address
+    def_delegators :@instance, :instance_id, :tags, :private_ip_address, :public_ip_address
 
     def start
       client = Aws::EC2::Client.new(client_options)
@@ -37,6 +37,11 @@ def start
       tag_instance
     end
 
+    def terminate
+      @instance.terminate
+      @instance.wait_until_terminated
+    end
+
     private
 
     def client_options

lib/ami_spec/aws_security_group.rb

@@ -0,0 +1,72 @@
+require 'aws-sdk-ec2'
+require 'forwardable'
+require 'securerandom'
+
+module AmiSpec
+  class AwsSecurityGroup
+    extend Forwardable
+
+    def self.create(**args)
+      new(**args).tap(&:create)
+    end
+
+    def initialize(ec2: Aws::EC2::Resource.new,
+                   group_name_prefix: "ami-spec-",
+                   connection_port: 22,
+                   subnet_id:,
+                   allow_any_ip: false,
+                   logger: Logger.new(STDOUT))
+      @ec2 = ec2
+      @group_name = "#{group_name_prefix}#{SecureRandom.uuid}"
+      @connection_port = connection_port
+      @subnet_id = subnet_id
+      @allow_any_ip = allow_any_ip
+      @logger = logger
+    end
+
+    def_delegators :@security_group, :group_id
+    attr_reader :group_name
+
+    def create
+      @logger.info "Creating temporary AWS security group: #{@group_name}"
+      create_security_group
+      allow_ingress_via_connection_port
+    end
+
+    def delete
+      @logger.info "Deleting temporary AWS security group: #{@group_name}"
+      @security_group.delete
+    end
+
+    private
+
+    def create_security_group
+      @security_group = @ec2.create_security_group(
+        group_name: @group_name,
+        description: "A temporary security group for running AmiSpec",
+        vpc_id: subnet.vpc_id,
+      )
+    end
+
+    def allow_ingress_via_connection_port
+      @security_group.authorize_ingress(
+        ip_permissions: [
+          {
+            ip_protocol: "tcp",
+            from_port: @connection_port,
+            to_port: @connection_port,
+            ip_ranges: [{cidr_ip: cidr_block}],
+          },
+        ],
+      )
+    end
+
+    def cidr_block
+      @allow_any_ip ? "0.0.0.0/0" : subnet.cidr_block
+    end
+
+    def subnet
+      @subnet ||= @ec2.subnet(@subnet_id)
+    end
+  end
+end

spec/ami_spec/aws_security_group_spec.rb

@@ -0,0 +1,92 @@
+require 'spec_helper'
+
+describe AmiSpec::AwsSecurityGroup do
+  subject(:aws_security_group) do
+    described_class.create(
+      ec2: ec2,
+      subnet_id: test_subnet_id,
+      allow_any_ip: allow_any_ip,
+      logger: logger
+    )
+  end
+
+  let(:ec2) { instance_spy(Aws::EC2::Resource, create_security_group: security_group, subnet: subnet) }
+  let(:security_group) { instance_spy(Aws::EC2::SecurityGroup, group_id: test_group_id) }
+  let(:subnet) { instance_spy(Aws::EC2::Subnet, vpc_id: test_vpc_id, cidr_block: test_cidr_block) }
+  let(:test_subnet_id) { 'test-subnet-id' }
+  let(:test_group_id) { 'test-group-id' }
+  let(:test_vpc_id) { 'test-vpc-id' }
+  let(:test_cidr_block) { '172.16.0.0/24' }
+  let(:allow_any_ip) { false }
+  let(:logger) { instance_spy(Logger) }
+
+  describe '#create' do
+    subject(:create) { aws_security_group }
+
+    it 'creates the security group in AWS' do
+      create
+      expect(ec2).to have_received(:create_security_group).with(group_name: aws_security_group.group_name, vpc_id: test_vpc_id, description: anything)
+    end
+
+    context 'given allow_any_ip: true' do
+      let(:allow_any_ip) { true }
+
+      it 'adds the ingress rule for SSH, allowing any IP address' do
+        create
+        expect(security_group).to have_received(:authorize_ingress).with(
+          ip_permissions: [{ip_protocol: "tcp", from_port: 22, to_port: 22, ip_ranges: [{cidr_ip: "0.0.0.0/0"}]}]
+        )
+      end
+    end
+
+    context 'given allow_any_ip: false' do
+      let(:allow_any_ip) { false }
+
+      it 'adds the ingress rule for SSH, allowing only IP addresses from the subnet CIDR block' do
+        create
+        expect(security_group).to have_received(:authorize_ingress).with(
+          ip_permissions: [{ip_protocol: "tcp", from_port: 22, to_port: 22, ip_ranges: [{cidr_ip: test_cidr_block}]}]
+        )
+      end
+    end
+
+    it 'loads the subnet to find the vpc id' do
+      create
+      expect(ec2).to have_received(:subnet).with(test_subnet_id)
+    end
+
+    it 'logs the creation of the key pair in AWS' do
+      create
+      expect(logger).to have_received(:info).with "Creating temporary AWS security group: #{aws_security_group.group_name}"
+    end
+  end
+
+  describe '#group_name' do
+    subject(:key_name) { aws_security_group.group_name }
+
+    it { should start_with('ami-spec-') }
+  end
+
+  describe '#group_id' do
+    subject(:group_id) { aws_security_group.group_id }
+
+    it 'is obtained from the AWS create-security-group API response' do
+      expect(group_id).to eq test_group_id
+    end
+  end
+
+  describe '#delete' do
+    subject(:delete) { aws_security_group.delete }
+
+    it 'deletes the security group in AWS' do
+      delete
+      expect(security_group).to have_received(:delete)
+    end
+
+    it 'logs the deletion of the key pair in AWS' do
+      delete
+      expect(logger).to have_received(:info).with "Deleting temporary AWS security group: #{aws_security_group.group_name}"
+    end
+  end
+end
+

spec/ami_spec_spec.rb

@@ -5,22 +5,28 @@
   let(:ec2) { instance_spy(Aws::EC2::Resource) }
   let(:ec2_double) { instance_double(AmiSpec::AwsInstance) }
   let(:aws_key_pair) { instance_spy(AmiSpec::AwsKeyPair) }
+  let(:aws_security_group) { instance_spy(AmiSpec::AwsSecurityGroup) }
   let(:state) { double(name: 'running') }
   let(:test_result) { true }
   let(:server_spec_double) { double(run: test_result) }
   let(:key_name) { 'key' }
+  let(:security_groups) { ['sg-1234'] }
+  let(:subnet_id) { 'subnet-1234abcd' }
+  let(:allow_any_temporary_security_group) { false }
   subject do
     described_class.run(
       amis: amis,
       specs: '/tmp/foobar',
-      subnet_id: 'subnet-1234abcd',
+      subnet_id: subnet_id,
       key_name: key_name,
       key_file: 'key.pem',
       aws_public_ip: false,
       aws_instance_type: 't2.micro',
       ssh_user: 'ubuntu',
       debug: false,
       ssh_retries: 30,
+      aws_security_groups: security_groups,
+      allow_any_temporary_security_group: allow_any_temporary_security_group,
     )
   end
 
@@ -37,6 +43,7 @@
       allow(AmiSpec::AwsInstance).to receive(:start).and_return(ec2_double)
       allow(AmiSpec::ServerSpec).to receive(:new).and_return(server_spec_double)
       allow(AmiSpec::AwsKeyPair).to receive(:create).and_return(aws_key_pair)
+      allow(AmiSpec::AwsSecurityGroup).to receive(:create).and_return(aws_security_group)
       allow(Aws::EC2::Resource).to receive(:new).and_return(ec2)
       allow(ec2_double).to receive(:terminate).and_return(true)
       allow(ec2_double).to receive(:private_ip_address).and_return('127.0.0.1')
@@ -85,6 +92,55 @@
         expect(AmiSpec::AwsKeyPair).not_to have_received(:create)
       end
     end
+
+    context 'given a security group id is not provided' do
+      let(:security_groups) { [] }
+
+      it 'creates a temporary security group' do
+        subject
+        expect(AmiSpec::AwsSecurityGroup).to have_received(:create)
+      end
+
+      it 'passes the subnet id' do
+        subject
+        expect(AmiSpec::AwsSecurityGroup).to have_received(:create)
+          .with(a_hash_including(subnet_id: subnet_id))
+      end
+
+      context 'given allow_any_temporary_security_group: true' do
+        let(:allow_any_temporary_security_group) { true }
+
+        it 'passes allow_any_ip: true' do
+          subject
+          expect(AmiSpec::AwsSecurityGroup).to have_received(:create)
+            .with(a_hash_including(allow_any_ip: true))
+        end
+      end
+
+      context 'given allow_any_temporary_security_group: false' do
+        let(:allow_any_temporary_security_group) { false }
+
+        it 'passes allow_any_ip: true' do
+          subject
+          expect(AmiSpec::AwsSecurityGroup).to have_received(:create)
+            .with(a_hash_including(allow_any_ip: false))
+        end
+      end
+
+      it 'deletes the temporary security group' do
+        subject
+        expect(aws_security_group).to have_received(:delete)
+      end
+    end
+
+    context 'given a security group id is provided' do
+      let(:security_groups) { ['sg-4321'] }
+
+      it 'does not create a temporary security group' do
+        subject
+        expect(AmiSpec::AwsSecurityGroup).not_to have_received(:create)
+      end
+    end
   end
 
   describe '#parse_tags' do

spec/aws_instance_spec.rb

@@ -6,9 +6,9 @@
   let(:role) { 'web_server' }
   let(:sec_group_id) { nil }
   let(:region) { nil }
-  let(:client_double) { instance_double(Aws::EC2::Client) }
-  let(:new_ec2_double) { instance_double(Aws::EC2::Types::Instance) }
-  let(:ec2_double) { instance_double(Aws::EC2::Instance) }
+  let(:client_double) { instance_spy(Aws::EC2::Client) }
+  let(:new_ec2_double) { instance_spy(Aws::EC2::Types::Instance) }
+  let(:ec2_double) { instance_spy(Aws::EC2::Instance) }
   let(:tags) { {} }
   let(:iam_instance_profile_arn) { nil }
   let(:user_data_file) { nil }
@@ -34,9 +34,6 @@
     allow(client_double).to receive(:run_instances).and_return(double(instances: [new_ec2_double]))
     allow(ec2_double).to receive(:create_tags).and_return(double)
     allow(Aws::EC2::Instance).to receive(:new).and_return(ec2_double)
-    allow(new_ec2_double).to receive(:instance_id)
-    allow(ec2_double).to receive(:instance_id)
-    allow(ec2_double).to receive(:wait_until_running)
   end
 
   describe '#start' do
@@ -139,4 +136,19 @@
     end
   end
 
+  describe '#terminate' do
+    subject(:terminate) { aws_instance.terminate }
+
+    before { aws_instance.start }
+
+    it 'instructs the EC2 instance to terminate' do
+      terminate
+      expect(ec2_double).to have_received(:terminate)
+    end
+
+    it 'waits for the EC2 instance to terminate' do
+      terminate
+      expect(ec2_double).to have_received(:wait_until_terminated)
+    end
+  end
 end