diff --git a/ebcli/controllers/ssh.py b/ebcli/controllers/ssh.py index f41dcfaec..4bc78ea75 100644 --- a/ebcli/controllers/ssh.py +++ b/ebcli/controllers/ssh.py @@ -35,6 +35,8 @@ class Meta: (['--setup'], dict( action='store_true', help=flag_text['ssh.setup'])), (['--timeout'], dict(type=int, help=flag_text['ssh.timeout'])), + (['--prefer-private-ip'], dict( + action='store_true', help=flag_text['ssh.prefer_private_ip'])), ] def do_command(self): @@ -47,6 +49,7 @@ def do_command(self): force = self.app.pargs.force setup = self.app.pargs.setup timeout = self.app.pargs.timeout + prefer_private_ip = self.app.pargs.prefer_private_ip if timeout and not setup: raise InvalidOptionsError(strings['ssh.timeout_without_setup']) @@ -60,5 +63,6 @@ def do_command(self): number=number, custom_ssh=custom_ssh, command=cmd, - timeout=timeout + timeout=timeout, + prefer_private_ip=prefer_private_ip, ) diff --git a/ebcli/operations/sshops.py b/ebcli/operations/sshops.py index 9a5e48f23..fd05c6854 100644 --- a/ebcli/operations/sshops.py +++ b/ebcli/operations/sshops.py @@ -28,7 +28,8 @@ def prepare_for_ssh(env_name, instance, keep_open, force, setup, number, keyname=None, no_keypair_error_message=None, - custom_ssh=None, command=None, timeout=None): + custom_ssh=None, command=None, timeout=None, + prefer_private_ip=False): if setup: setup_ssh(env_name, keyname, timeout=timeout) return @@ -62,7 +63,8 @@ def prepare_for_ssh(env_name, instance, keep_open, force, setup, number, keep_open=keep_open, force_open=force, custom_ssh=custom_ssh, - command=command + command=command, + prefer_private_ip=prefer_private_ip, ) except NoKeypairError: if not no_keypair_error_message: @@ -85,19 +87,20 @@ def setup_ssh(env_name, keyname, timeout=None): commonops.update_environment(env_name, options, False, timeout=timeout or 5) -def ssh_into_instance(instance_id, keep_open=False, force_open=False, custom_ssh=None, command=None): +def ssh_into_instance(instance_id, keep_open=False, force_open=False, custom_ssh=None, command=None, prefer_private_ip=False): instance = ec2.describe_instance(instance_id) try: keypair_name = instance['KeyName'] except KeyError: raise NoKeypairError() - try: - ip = instance['PublicIpAddress'] - except KeyError: - if 'PrivateIpAddress' in instance: - ip = instance['PrivateIpAddress'] - else: - raise NotFoundError(strings['ssh.noip']) + + if prefer_private_ip: + ip = instance.get('PrivateIpAddress', instance.get('PublicIpAddress')) + else: + ip = instance.get('PublicIpAddress', instance.get('PrivateIpAddress')) + if ip is None: + raise NotFoundError(strings['ssh.noip']) + security_groups = instance['SecurityGroups'] user = 'ec2-user' diff --git a/ebcli/resources/strings.py b/ebcli/resources/strings.py index 59cda8ad6..c72e8786c 100644 --- a/ebcli/resources/strings.py +++ b/ebcli/resources/strings.py @@ -874,6 +874,7 @@ 'ssh.setup': 'setup SSH for the environment', 'ssh.timeout': "Specify the timeout period in minutes. Can only be used with the " "'--setup' argument.", + 'ssh.prefer_private_ip': "Prefer connecting to an instance's private IP address", 'cleanup.resources': 'Valid values include (builder, versions, all). You can specify ' '"builder" to terminate the environment used to create this platform. ' diff --git a/tests/unit/controllers/test_ssh.py b/tests/unit/controllers/test_ssh.py index 53bdac297..6faa0ec94 100644 --- a/tests/unit/controllers/test_ssh.py +++ b/tests/unit/controllers/test_ssh.py @@ -58,7 +58,8 @@ def test_ssh( keep_open=False, number=None, setup=False, - timeout=None + timeout=None, + prefer_private_ip=False ) @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name') @@ -83,7 +84,8 @@ def test_ssh__setup__with_timeout( keep_open=False, number=None, setup=True, - timeout=10 + timeout=10, + prefer_private_ip=False ) @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name') @@ -108,7 +110,8 @@ def test_ssh__setup__without_timeout( keep_open=False, number=None, setup=True, - timeout=None + timeout=None, + prefer_private_ip=False ) @mock.patch('ebcli.controllers.ssh.SSHController.get_env_name') diff --git a/tests/unit/operations/test_sshops.py b/tests/unit/operations/test_sshops.py index 7047236d2..73b68a57c 100644 --- a/tests/unit/operations/test_sshops.py +++ b/tests/unit/operations/test_sshops.py @@ -321,6 +321,28 @@ def test_ssh_into_instance__uses_private_address( sshops.ssh_into_instance('instance-id') call_mock.assert_called_once_with(['ssh', '-i', 'aws-eb-us-west-2', 'ec2-user@172.31.35.210']) + + @mock.patch('ebcli.operations.sshops.ec2.describe_instance') + @mock.patch('ebcli.operations.sshops.ec2.describe_security_group') + @mock.patch('ebcli.operations.sshops.ec2.authorize_ssh') + @mock.patch('ebcli.operations.sshops._get_ssh_file') + @mock.patch('ebcli.operations.sshops.subprocess.call') + def test_ssh_into_instance__uses_private_address_when_private_ip_flag_is_present( + self, + call_mock, + _get_ssh_file_mock, + authorize_ssh_mock, + describe_security_group_mock, + describe_instance_mock + ): + describe_instance_response = deepcopy(mock_responses.DESCRIBE_INSTANCES_RESPONSE['Reservations'][0]['Instances'][0]) + describe_instance_mock.return_value = describe_instance_response + describe_security_group_mock.return_value = mock_responses.DESCRIBE_SECURITY_GROUPS_RESPONSE['SecurityGroups'][0] + _get_ssh_file_mock.return_value = 'aws-eb-us-west-2' + call_mock.return_value = 0 + + sshops.ssh_into_instance('instance-id', prefer_private_ip=True) + call_mock.assert_called_once_with(['ssh', '-i', 'aws-eb-us-west-2', 'ec2-user@172.31.35.210']) @mock.patch('ebcli.operations.sshops.ec2.describe_instance') @mock.patch('ebcli.operations.sshops.ec2.describe_security_group') @@ -520,7 +542,8 @@ def test_prepare_for_ssh__choose_instance_to_ssh_into( command=None, custom_ssh=None, force_open=False, - keep_open=False + keep_open=False, + prefer_private_ip=False ) @mock.patch('ebcli.operations.sshops.commonops.get_instance_ids') @@ -553,7 +576,8 @@ def test_prepare_for_ssh__number_of_instance_specified( command=None, custom_ssh=None, force_open=False, - keep_open=False + keep_open=False, + prefer_private_ip=False ) @mock.patch('ebcli.operations.sshops.commonops.get_instance_ids') @@ -589,7 +613,8 @@ def test_prepare_for_ssh__ssh_into_instance_fails( command=None, custom_ssh=None, force_open=False, - keep_open=False + keep_open=False, + prefer_private_ip=False ) log_error_mock.assert_called_once_with( 'This environment is not set up for SSH. Use "eb ssh --setup" to set up SSH for the environment.'