def test_ssh_public_key(self, key_type):
key = SshKey(key_type=key_type)
key.generate()
if key_type == 'rsa':
key_start_str = 'ssh-rsa'
else:
key_start_str = 'ssh-ed25519'
self.assertTrue(key.public_key.decode().startswith(key_start_str))
def test_ssh_public_key(key_type):
key = SshKey(key_type=key_type)
key.generate()
if key_type == "rsa":
key_start_str = "ssh-rsa"
else:
key_start_str = "ssh-ed25519"
assert key.public_key.decode().startswith(key_start_str)
def test_initialize_key_as_context_manager(mocker, key_type):
mocker.patch("aws_gate.ssh_common.os")
open_mock = mocker.patch("builtins.open", new_callable=mocker.mock_open())
with SshKey(key_type=key_type):
assert open_mock.called
open_mock.assert_called_with(DEFAULT_GATE_KEY_PATH, "wb")
def test_delete_key(self):
with patch('builtins.open', new_callable=mock_open()), \
patch('aws_gate.ssh_common.os', new_callable=MagicMock()) as m:
key = SshKey()
key.generate()
key.write_to_file()
key.delete()
self.assertTrue(m.remove.called)
self.assertEqual(m.remove.call_args, call(DEFAULT_GATE_KEY_PATH))
def test_delete_key(mocker):
mocker.patch("builtins.open", new_callable=mocker.mock_open())
m = mocker.patch("aws_gate.ssh_common.os", new_callable=mocker.MagicMock())
key = SshKey()
key.generate()
key.write_to_file()
key.delete()
assert m.remove.called
assert m.remove.call_args == mocker.call(DEFAULT_GATE_KEY_PATH)
def ssh(
config,
instance_name,
user=DEFAULT_OS_USER,
port=DEFAULT_SSH_PORT,
key_type=DEFAULT_KEY_ALGORITHM,
key_size=DEFAULT_KEY_SIZE,
profile_name=AWS_DEFAULT_PROFILE,
region_name=AWS_DEFAULT_REGION,
command=None,
forwarding=None,
):
instance, profile, region = fetch_instance_details_from_config(
config, instance_name, profile_name, region_name)
ssm = get_aws_client("ssm", region_name=region, profile_name=profile)
ec2 = get_aws_resource("ec2", region_name=region, profile_name=profile)
ec2_ic = get_aws_client("ec2-instance-connect",
region_name=region,
profile_name=profile)
instance_id = query_instance(name=instance, ec2=ec2)
if instance_id is None:
raise ValueError(
"No instance could be found for name: {}".format(instance))
az = get_instance_details(instance_id=instance_id,
ec2=ec2)["availability_zone"]
logger.info(
"Opening SSH session on instance %s (%s) via profile %s",
instance_id,
region,
profile,
)
with SshKey(key_type=key_type, key_size=key_size) as ssh_key:
with SshKeyUploader(instance_id=instance_id,
az=az,
user=user,
ssh_key=ssh_key,
ec2_ic=ec2_ic):
with SshSession(
instance_id,
region_name=region,
profile_name=profile,
ssm=ssm,
port=port,
user=user,
command=command,
forwarding=forwarding,
) as ssh_session:
ssh_session.open()
def test_ssh_key_file_permissions(self):
with patch('builtins.open', new_callable=mock_open()), \
patch('aws_gate.ssh_common.os.chmod') as m:
key = SshKey()
key.generate()
key.write_to_file()
self.assertTrue(m.called)
self.assertEqual(call(DEFAULT_GATE_KEY_PATH, 0o600), m.call_args)
def test_ssh_key_file_permissions(mocker):
mocker.patch("builtins.open", new_callable=mocker.mock_open())
m = mocker.patch("aws_gate.ssh_common.os.chmod")
key = SshKey()
key.generate()
key.write_to_file()
assert m.called
assert mocker.call(DEFAULT_GATE_KEY_PATH, 0o600) == m.call_args_list[0]
def ssh_proxy(config, instance_name, user=DEFAULT_OS_USER, port=DEFAULT_SSH_PORT, key_type=DEFAULT_KEY_ALGORITHM,
key_size=DEFAULT_KEY_SIZE, profile_name=AWS_DEFAULT_PROFILE, region_name=AWS_DEFAULT_REGION):
instance, profile, region = fetch_instance_details_from_config(config, instance_name, profile_name, region_name)
ssm = get_aws_client('ssm', region_name=region, profile_name=profile)
ec2 = get_aws_resource('ec2', region_name=region, profile_name=profile)
ec2_ic = get_aws_client('ec2-instance-connect', region_name=region, profile_name=profile)
instance_id = query_instance(name=instance, ec2=ec2)
if instance_id is None:
raise ValueError('No instance could be found for name: {}'.format(instance))
az = get_instance_details(instance_id=instance_id, ec2=ec2)['availability_zone']
logger.info('Opening SSH proxy session on instance %s (%s) via profile %s', instance_id, region, profile)
with SshKey(key_type=key_type, key_size=key_size) as ssh_key:
with SshKeyUploader(instance_id=instance_id, az=az, user=user, ssh_key=ssh_key, ec2_ic=ec2_ic):
with SshProxySession(instance_id, region_name=region, profile_name=profile, ssm=ssm, port=port,
user=user) as ssh_proxy_session:
ssh_proxy_session.open()
def test_initialize_key_unsupported_key_type(key_type):
with pytest.raises(ValueError):
SshKey(key_type=key_type)
def test_initialize_key_invalid_key_path():
with pytest.raises(ValueError):
SshKey(key_path="")
def test_initialize_key(key_type, key_size):
key = SshKey(key_type=key_type, key_size=key_size)
assert key.key_path == DEFAULT_GATE_KEY_PATH
assert key.key_type == key_type
assert key.key_size == key_size
def test_initialize_key_as_context_manager(self, key_type):
with patch('builtins.open', new_callable=mock_open()) as open_mock, \
patch('aws_gate.ssh_common.os'):
with SshKey(key_type=key_type):
self.assertTrue(open_mock.called)
open_mock.assert_called_with(DEFAULT_GATE_KEY_PATH, 'wb')
def test_initialize_key_invalid_key_path(self):
with self.assertRaises(ValueError):
SshKey(key_path='')
def test_initialize_key_unsupported_key_size(self, key_size):
with self.assertRaises(ValueError):
SshKey(key_size=key_size)
def test_initialize_key(self, key_type, key_size):
key = SshKey(key_type=key_type)
self.assertTrue(key.key_path, DEFAULT_GATE_KEY_PATH)
self.assertTrue(key.key_type, key_type)
self.assertTrue(key.key_size, key_size)