def test_subnet_with_security_groups(self):
""" Test that the subnet and security groups are passed to the
calls to AWSService.run_instance().
"""
self.call_count = 0
def run_instance_callback(args):
self.call_count += 1
self.assertEqual('subnet-1', args.subnet_id)
if self.call_count == 1:
# Snapshotter.
self.assertIsNone(args.security_group_ids)
elif self.call_count == 2:
# Encryptor.
self.assertEqual(['sg-1', 'sg-2'], args.security_group_ids)
else:
self.fail('Unexpected number of calls to run_instance()')
aws_svc, encryptor_image, guest_image = build_aws_service()
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
subnet_id='subnet-1',
security_group_ids=['sg-1', 'sg-2']
)
def test_delete_orphaned_volumes(self):
""" Test that we clean up instance volumes that are orphaned by AWS.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
# Simulate a tagged orphaned volume.
volume = Volume()
volume.id = test_aws_service.new_id()
aws_svc.volumes[volume.id] = volume
aws_svc.tagged_volumes.append(volume)
# Verify that lookup succeeds before encrypt().
self.assertEqual(volume, aws_svc.get_volume(volume.id))
self.assertEqual([volume],
aws_svc.get_volumes(
tag_key=encrypt_ami.TAG_ENCRYPTOR_SESSION_ID,
tag_value='123'))
encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
# Verify that the volume was deleted.
self.assertIsNone(aws_svc.volumes.get(volume.id, None))
def test_subnet_without_security_groups(self):
""" Test that we create the temporary security group in the subnet
that the user specified.
"""
self.security_group_was_created = False
def create_security_group_callback(vpc_id):
self.security_group_was_created = True
self.assertEqual('vpc-1', vpc_id)
aws_svc, encryptor_image, guest_image = build_aws_service()
aws_svc.create_security_group_callback = \
create_security_group_callback
subnet = Subnet()
subnet.id = 'subnet-1'
subnet.vpc_id = 'vpc-1'
aws_svc.subnets = {subnet.id: subnet}
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
subnet_id='subnet-1'
)
self.assertTrue(self.security_group_was_created)
def test_terminate_guest(self):
""" Test that we terminate the guest instance if an exception is
raised while waiting for it to come up.
"""
self.terminate_instance_called = False
self.instance_id = None
def get_instance_callback(instance):
self.instance_id = instance.id
raise TestException('Test')
def terminate_instance_callback(instance):
self.terminate_instance_called = True
self.assertEqual(self.instance_id, instance.id)
aws_svc, encryptor_image, guest_image = build_aws_service()
aws_svc.get_instance_callback = get_instance_callback
aws_svc.terminate_instance_callback = terminate_instance_callback
try:
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
except TestException:
pass
self.assertTrue(self.terminate_instance_called)
def test_security_group_eventual_consistency(self):
""" Test that we handle eventually consistency issues when creating
a temporary security group.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
self.call_count = 0
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.call_count += 1
if self.call_count < 3:
# Simulate eventual consistency error while creating
# security group.
e = EC2ResponseError(None, None)
e.error_code = 'InvalidGroup.NotFound'
raise e
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.assertEquals(3, self.call_count)
def test_clean_up_encryptor_instance(self):
""" Test that we clean up the encryptor instance if an exception is
raised inside _run_encryptor_instance().
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
self.encryptor_instance_id = None
self.encryptor_terminated = False
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.encryptor_instance_id = args.instance.id
def create_tags_callback(resource_id, name, description):
if resource_id == self.encryptor_instance_id:
raise TestException()
def terminate_instance_callback(instance_id):
if instance_id == self.encryptor_instance_id:
self.encryptor_terminated = True
aws_svc.run_instance_callback = run_instance_callback
aws_svc.create_tags_callback = create_tags_callback
aws_svc.terminate_instance_callback = terminate_instance_callback
with self.assertRaises(TestException):
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.assertTrue(self.encryptor_terminated)
def test_delete_orphaned_volumes(self):
""" Test that we clean up instance volumes that are orphaned by AWS.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
# Simulate a tagged orphaned volume.
volume = Volume()
volume.id = test_aws_service.new_id()
aws_svc.volumes[volume.id] = volume
aws_svc.tagged_volumes.append(volume)
# Verify that lookup succeeds before encrypt().
self.assertEqual(volume, aws_svc.get_volume(volume.id))
self.assertEqual(
[volume],
aws_svc.get_volumes(
tag_key=encrypt_ami.TAG_ENCRYPTOR_SESSION_ID, tag_value='123')
)
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
# Verify that the volume was deleted.
self.assertIsNone(aws_svc.volumes.get(volume.id, None))
def test_brkt_env_encrypt(self):
""" Test that we parse the brkt_env value and pass the correct
values to user_data when launching the encryptor instance.
"""
api_host_port = 'api.example.com:777'
hsmproxy_host_port = 'hsmproxy.example.com:888'
aws_svc, encryptor_image, guest_image = build_aws_service()
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
brkt_config = self._get_brkt_config_from_mime(args.user_data)
d = json.loads(brkt_config)
self.assertEquals(
api_host_port,
d['brkt']['api_host']
)
self.assertEquals(
hsmproxy_host_port,
d['brkt']['hsmproxy_host']
)
cli_args = '--brkt-env %s,%s' % (api_host_port, hsmproxy_host_port)
values = instance_config_args_to_values(cli_args)
brkt_env = brkt_cli.brkt_env_from_values(values)
ic = make_instance_config(values, brkt_env)
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
instance_config=ic
)
def test_encryption_error_console_output_not_available(self):
""" Test that we handle the case when encryption fails and console
output is not available.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
aws_svc.console_output_text = None
try:
encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=FailedEncryptionService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
self.fail('Encryption should have failed')
except encryptor_service.EncryptionError as e:
self.assertIsNone(e.console_output_file)
def test_guest_instance_type(self):
""" Test that the guest instance type is passed through
to run_instance().
"""
aws_svc, encryptor_image, guest_image = \
test_aws_service.build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
def run_instance_callback(args):
if args.image_id == encrypted_ami_id:
self.assertEqual('t2.micro', args.instance_type)
elif args.image_id == encryptor_image.id:
self.assertEqual('m3.medium', args.instance_type)
else:
self.fail('Unexpected image: ' + args.image_id)
aws_svc.run_instance_callback = run_instance_callback
update_ami(
aws_svc, encrypted_ami_id, encryptor_image.id, 'Test updated AMI',
subnet_id='subnet-1', security_group_ids=['sg-1', 'sg-2'],
enc_svc_class=DummyEncryptorService, guest_instance_type='t2.micro'
)
def test_brkt_env_update(self):
""" Test that the Bracket environment is passed through to metavisor
user data.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
api_host_port = 'api.example.com:777'
hsmproxy_host_port = 'hsmproxy.example.com:888'
cli_args = '--brkt-env %s,%s' % (api_host_port, hsmproxy_host_port)
values = instance_config_args_to_values(cli_args)
brkt_env = brkt_cli.brkt_env_from_values(values)
ic = make_instance_config(values, brkt_env)
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
brkt_config = self._get_brkt_config_from_mime(args.user_data)
d = json.loads(brkt_config)
self.assertEquals(api_host_port, d['brkt']['api_host'])
self.assertEquals(hsmproxy_host_port,
d['brkt']['hsmproxy_host'])
self.assertEquals('updater', d['brkt']['solo_mode'])
aws_svc.run_instance_callback = run_instance_callback
update_ami(aws_svc,
encrypted_ami_id,
encryptor_image.id,
'Test updated AMI',
enc_svc_class=DummyEncryptorService,
instance_config=ic)
def test_guest_instance_type(self):
""" Test that we use the specified instance type to launch the guest
instance.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
def run_instance_callback(args):
if args.image_id == guest_image.id:
self.assertEqual('t2.micro', args.instance_type)
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
guest_instance_type='t2.micro')
def test_subnet_and_security_groups(self):
""" Test that the subnet and security group ids are passed through
to run_instance().
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
self.call_count = 0
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.call_count += 1
self.assertEqual('subnet-1', args.subnet_id)
self.assertEqual(['sg-1', 'sg-2'], args.security_group_ids)
aws_svc.run_instance_callback = run_instance_callback
ami_id = update_ami(aws_svc,
encrypted_ami_id,
encryptor_image.id,
'Test updated AMI',
subnet_id='subnet-1',
security_group_ids=['sg-1', 'sg-2'],
enc_svc_class=DummyEncryptorService)
self.assertEqual(1, self.call_count)
self.assertIsNotNone(ami_id)
def test_guest_instance_type(self):
""" Test that the guest instance type is passed through
to run_instance().
"""
aws_svc, encryptor_image, guest_image = \
test_aws_service.build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
def run_instance_callback(args):
if args.image_id == encrypted_ami_id:
self.assertEqual('t2.micro', args.instance_type)
elif args.image_id == encryptor_image.id:
self.assertEqual('m3.medium', args.instance_type)
else:
self.fail('Unexpected image: ' + args.image_id)
aws_svc.run_instance_callback = run_instance_callback
update_ami(aws_svc,
encrypted_ami_id,
encryptor_image.id,
'Test updated AMI',
subnet_id='subnet-1',
security_group_ids=['sg-1', 'sg-2'],
enc_svc_class=DummyEncryptorService,
guest_instance_type='t2.micro')
def test_subnet_and_security_groups(self):
""" Test that the subnet and security group ids are passed through
to run_instance().
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.call_count = 0
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.call_count += 1
self.assertEqual('subnet-1', args.subnet_id)
self.assertEqual(['sg-1', 'sg-2'], args.security_group_ids)
aws_svc.run_instance_callback = run_instance_callback
ami_id = update_ami(
aws_svc, encrypted_ami_id, encryptor_image.id,
'Test updated AMI',
subnet_id='subnet-1', security_group_ids=['sg-1', 'sg-2'],
enc_svc_class=DummyEncryptorService
)
self.assertEqual(1, self.call_count)
self.assertIsNotNone(ami_id)
def test_security_group_eventual_consistency(self):
""" Test that we handle eventually consistency issues when creating
a temporary security group.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
self.call_count = 0
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.call_count += 1
if self.call_count < 3:
# Simulate eventual consistency error while creating
# security group.
e = EC2ResponseError(None, None)
e.error_code = 'InvalidGroup.NotFound'
raise e
aws_svc.run_instance_callback = run_instance_callback
update_ami(aws_svc,
encrypted_ami_id,
encryptor_image.id,
'Test updated AMI',
enc_svc_class=DummyEncryptorService)
self.assertEqual(3, self.call_count)
def test_encryption_error_console_output_not_available(self):
""" Test that we handle the case when encryption fails and console
output is not available.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
aws_svc.console_output_text = None
try:
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=FailedEncryptionService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.fail('Encryption should have failed')
except encryptor_service.EncryptionError as e:
self.assertIsNone(e.console_output_file)
def test_guest_instance_type(self):
""" Test that we use the specified instance type to launch the guest
instance.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
def run_instance_callback(args):
if args.image_id == guest_image.id:
self.assertEqual('t2.micro', args.instance_type)
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
guest_instance_type='t2.micro'
)
def test_smoke(self):
""" Run the entire process and test that nothing obvious is broken.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
self.assertIsNotNone(encrypted_ami_id)
def command_encrypt_ami(values):
session_id = util.make_nonce()
aws_svc = aws_service.AWSService(
session_id,
retry_timeout=values.retry_timeout,
retry_initial_sleep_seconds=values.retry_initial_sleep_seconds)
log.debug('Retry timeout=%.02f, initial sleep seconds=%.02f',
aws_svc.retry_timeout, aws_svc.retry_initial_sleep_seconds)
brkt_env = (brkt_cli.brkt_env_from_values(values)
or brkt_cli.get_prod_brkt_env())
if values.validate:
# Validate the region before connecting.
_validate_region(aws_svc, values.region)
if values.token:
brkt_cli.check_jwt_auth(brkt_env, values.token)
aws_svc.connect(values.region, key_name=values.key_name)
if values.validate:
guest_image = _validate_guest_ami(aws_svc, values.ami)
else:
guest_image = aws_svc.get_image(values.ami)
pv = _use_pv_metavisor(values, guest_image)
encryptor_ami = (values.encryptor_ami
or _get_encryptor_ami(values.region, pv=pv))
default_tags = encrypt_ami.get_default_tags(session_id, encryptor_ami)
default_tags.update(brkt_cli.parse_tags(values.tags))
aws_svc.default_tags = default_tags
if values.validate:
_validate(aws_svc, values, encryptor_ami)
brkt_cli.validate_ntp_servers(values.ntp_servers)
encrypted_image_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=encryptor_service.EncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_ami,
encrypted_ami_name=values.encrypted_ami_name,
subnet_id=values.subnet_id,
security_group_ids=values.security_group_ids,
guest_instance_type=values.guest_instance_type,
instance_config=make_instance_config(values, brkt_env),
status_port=values.status_port,
save_encryptor_logs=values.save_encryptor_logs)
# Print the AMI ID to stdout, in case the caller wants to process
# the output. Log messages go to stderr.
print(encrypted_image_id)
return 0
def test_smoke(self):
""" Run the entire process and test that nothing obvious is broken.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.assertIsNotNone(encrypted_ami_id)
def test_instance_type(self):
""" Test that we launch the guest as m3.medium and the encryptor
as c3.xlarge.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
def run_instance_callback(args):
if args.image_id == guest_image.id:
self.assertEqual('m3.medium', args.instance_type)
self.assertFalse(args.ebs_optimized)
elif args.image_id == encryptor_image.id:
self.assertEqual('c3.xlarge', args.instance_type)
self.assertTrue(args.ebs_optimized)
else:
self.fail('Unexpected image id: ' + args.image_id)
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
def test_encrypted_ami_name(self):
""" Test that the name is set on the encrypted AMI when specified.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
name = 'Am I an AMI?'
image_id = encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
encrypted_ami_name=name)
ami = aws_svc.get_image(image_id)
self.assertEqual(name, ami.name)
def test_encryption_error_console_output_available(self):
""" Test that when an encryption failure occurs, we write the
console log to a temp file.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
# Create callbacks that make sure that we stop the encryptor
# instance before collecting logs.
self.encryptor_instance = None
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.encryptor_instance = args.instance
self.encryptor_stopped = False
def stop_instance_callback(instance):
if (self.encryptor_instance and
instance.id == self.encryptor_instance.id):
self.encryptor_stopped = True
aws_svc.run_instance_callback = run_instance_callback
aws_svc.stop_instance_callback = stop_instance_callback
try:
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=FailedEncryptionService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
self.fail('Encryption should have failed')
except encryptor_service.EncryptionError as e:
with open(e.console_output_file.name) as f:
content = f.read()
self.assertEquals(test_aws_service.CONSOLE_OUTPUT_TEXT, content)
os.remove(e.console_output_file.name)
self.assertTrue(self.encryptor_stopped)
def test_encryption_error_console_output_available(self):
""" Test that when an encryption failure occurs, we write the
console log to a temp file.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
# Create callbacks that make sure that we stop the encryptor
# instance before collecting logs.
self.encryptor_instance = None
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.encryptor_instance = args.instance
self.encryptor_stopped = False
def stop_instance_callback(instance):
if (self.encryptor_instance
and instance.id == self.encryptor_instance.id):
self.encryptor_stopped = True
aws_svc.run_instance_callback = run_instance_callback
aws_svc.stop_instance_callback = stop_instance_callback
try:
encrypt_ami.encrypt(aws_svc=aws_svc,
enc_svc_cls=FailedEncryptionService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id)
self.fail('Encryption should have failed')
except encryptor_service.EncryptionError as e:
with open(e.console_output_file.name) as f:
content = f.read()
self.assertEquals(test_aws_service.CONSOLE_OUTPUT_TEXT,
content)
os.remove(e.console_output_file.name)
self.assertTrue(self.encryptor_stopped)
def test_instance_type(self):
""" Test that we launch the guest as m3.medium and the encryptor
as c3.xlarge.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
def run_instance_callback(args):
if args.image_id == guest_image.id:
self.assertEqual('m3.medium', args.instance_type)
self.assertFalse(args.ebs_optimized)
elif args.image_id == encryptor_image.id:
self.assertEqual('c3.xlarge', args.instance_type)
self.assertTrue(args.ebs_optimized)
else:
self.fail('Unexpected image id: ' + args.image_id)
aws_svc.run_instance_callback = run_instance_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
def test_encrypted_ami_name(self):
""" Test that the name is set on the encrypted AMI when specified.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
name = 'Am I an AMI?'
image_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
encrypted_ami_name=name
)
ami = aws_svc.get_image(image_id)
self.assertEqual(name, ami.name)
def test_brkt_env_update(self):
""" Test that the Bracket environment is passed through to metavisor
user data.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
encrypted_ami_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id
)
api_host_port = 'api.example.com:777'
hsmproxy_host_port = 'hsmproxy.example.com:888'
network_host_port = 'network.example.com:999'
cli_args = '--brkt-env %s,%s,%s' % (api_host_port, hsmproxy_host_port,
network_host_port)
values = instance_config_args_to_values(cli_args)
ic = instance_config_from_values(values)
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
brkt_config = self._get_brkt_config_from_mime(args.user_data)
d = json.loads(brkt_config)
self.assertEquals(
api_host_port,
d['brkt']['api_host']
)
self.assertEquals(
hsmproxy_host_port,
d['brkt']['hsmproxy_host']
)
self.assertEquals(
network_host_port,
d['brkt']['network_host']
)
self.assertEquals(
'updater',
d['brkt']['solo_mode']
)
aws_svc.run_instance_callback = run_instance_callback
update_ami(
aws_svc, encrypted_ami_id, encryptor_image.id,
'Test updated AMI',
enc_svc_class=DummyEncryptorService,
instance_config=ic
)
def test_no_terminate_encryptor_on_failure(self):
""" Test that when terminate_encryptor_on_failure=False, we terminate
the encryptor only when encryption succeeds.
"""
aws_svc, encryptor_image, guest_image = build_aws_service()
self.guest_instance_id = None
self.guest_terminated = False
self.encryptor_instance_id = None
self.encryptor_terminated = False
self.security_group_deleted = False
self.snapshot_deleted = False
def run_instance_callback(args):
if args.image_id == encryptor_image.id:
self.encryptor_instance_id = args.instance.id
if args.image_id == guest_image.id:
self.guest_instance_id = args.instance.id
def delete_snapshot_callback(snapshot_id):
self.snapshot_deleted = True
def terminate_instance_callback(instance_id):
self.assertIsNotNone(self.encryptor_instance_id)
if instance_id == self.encryptor_instance_id:
self.encryptor_terminated = True
if instance_id == self.guest_instance_id:
self.guest_terminated = True
def delete_security_group_callback(sg_id):
self.security_group_deleted = True
# Encryption succeeded. Make sure the encryptor was terminated.
aws_svc.run_instance_callback = run_instance_callback
aws_svc.delete_snapshot_callback = delete_snapshot_callback
aws_svc.terminate_instance_callback = terminate_instance_callback
aws_svc.delete_security_group_callback = \
delete_security_group_callback
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=DummyEncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
terminate_encryptor_on_failure=False
)
self.assertIsNotNone(self.encryptor_instance_id)
self.assertTrue(self.snapshot_deleted)
self.assertTrue(self.guest_terminated)
self.assertTrue(self.encryptor_terminated)
self.assertTrue(self.security_group_deleted)
# Encryption failed. Make sure the encryptor was not terminated.
self.guest_instance_id = None
self.guest_terminated = False
self.encryptor_instance_id = None
self.encryptor_terminated = False
self.security_group_deleted = False
self.snapshot_deleted = False
with self.assertRaises(encryptor_service.EncryptionError):
encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=FailedEncryptionService,
image_id=guest_image.id,
encryptor_ami=encryptor_image.id,
terminate_encryptor_on_failure=False
)
self.assertIsNotNone(self.encryptor_instance_id)
self.assertTrue(self.snapshot_deleted)
self.assertTrue(self.guest_terminated)
self.assertFalse(self.encryptor_terminated)
self.assertFalse(self.security_group_deleted)
def run_encrypt(values, config, verbose=False):
session_id = util.make_nonce()
aws_svc = aws_service.AWSService(
session_id,
retry_timeout=values.retry_timeout,
retry_initial_sleep_seconds=values.retry_initial_sleep_seconds
)
log.debug(
'Retry timeout=%.02f, initial sleep seconds=%.02f',
aws_svc.retry_timeout, aws_svc.retry_initial_sleep_seconds)
brkt_env = (
brkt_cli.brkt_env_from_values(values) or
brkt_cli.get_prod_brkt_env()
)
if values.validate:
# Validate the region before connecting.
_validate_region(aws_svc, values.region)
if values.token:
brkt_cli.check_jwt_auth(brkt_env, values.token)
aws_svc.connect(values.region, key_name=values.key_name)
if values.validate:
guest_image = _validate_guest_ami(aws_svc, values.ami)
else:
guest_image = aws_svc.get_image(values.ami)
encryptor_ami = values.encryptor_ami or _get_encryptor_ami(values.region)
default_tags = encrypt_ami.get_default_tags(session_id, encryptor_ami)
default_tags.update(brkt_cli.parse_tags(values.tags))
aws_svc.default_tags = default_tags
if values.validate:
_validate(aws_svc, values, encryptor_ami)
brkt_cli.validate_ntp_servers(values.ntp_servers)
instance_config = instance_config_from_values(
values, mode=INSTANCE_CREATOR_MODE, cli_config=config)
if verbose:
with tempfile.NamedTemporaryFile(
prefix='user-data-',
delete=False
) as f:
log.debug('Writing instance user data to %s', f.name)
f.write(instance_config.make_userdata())
encrypted_image_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=encryptor_service.EncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_ami,
encrypted_ami_name=values.encrypted_ami_name,
subnet_id=values.subnet_id,
security_group_ids=values.security_group_ids,
guest_instance_type=values.guest_instance_type,
instance_config=instance_config,
status_port=values.status_port,
save_encryptor_logs=values.save_encryptor_logs,
terminate_encryptor_on_failure=(
values.terminate_encryptor_on_failure)
)
# Print the AMI ID to stdout, in case the caller wants to process
# the output. Log messages go to stderr.
print(encrypted_image_id)
return 0
def command_encrypt_ami(values):
session_id = util.make_nonce()
aws_svc = aws_service.AWSService(
session_id,
retry_timeout=values.retry_timeout,
retry_initial_sleep_seconds=values.retry_initial_sleep_seconds
)
log.debug(
'Retry timeout=%.02f, initial sleep seconds=%.02f',
aws_svc.retry_timeout, aws_svc.retry_initial_sleep_seconds)
brkt_env = (
brkt_cli.brkt_env_from_values(values) or
brkt_cli.get_prod_brkt_env()
)
if values.validate:
# Validate the region before connecting.
_validate_region(aws_svc, values.region)
if values.token:
brkt_cli.check_jwt_auth(brkt_env, values.token)
aws_svc.connect(values.region, key_name=values.key_name)
if values.validate:
guest_image = _validate_guest_ami(aws_svc, values.ami)
else:
guest_image = aws_svc.get_image(values.ami)
pv = _use_pv_metavisor(values, guest_image)
encryptor_ami = (
values.encryptor_ami or
_get_encryptor_ami(values.region, pv=pv)
)
default_tags = encrypt_ami.get_default_tags(session_id, encryptor_ami)
default_tags.update(brkt_cli.parse_tags(values.tags))
aws_svc.default_tags = default_tags
if values.validate:
_validate(aws_svc, values, encryptor_ami)
brkt_cli.validate_ntp_servers(values.ntp_servers)
encrypted_image_id = encrypt_ami.encrypt(
aws_svc=aws_svc,
enc_svc_cls=encryptor_service.EncryptorService,
image_id=guest_image.id,
encryptor_ami=encryptor_ami,
encrypted_ami_name=values.encrypted_ami_name,
subnet_id=values.subnet_id,
security_group_ids=values.security_group_ids,
guest_instance_type=values.guest_instance_type,
instance_config=make_instance_config(values, brkt_env),
status_port=values.status_port,
save_encryptor_logs=values.save_encryptor_logs
)
# Print the AMI ID to stdout, in case the caller wants to process
# the output. Log messages go to stderr.
print(encrypted_image_id)
return 0
