Merge pull request #35 from zalando-stups/validate-instance

Validate instance

piu/cli.py

@@ -3,7 +3,6 @@
 Helper script to request access to a certain host.
 '''
 
-import boto3
 import click
 import datetime
 import operator
@@ -142,6 +141,12 @@ def tunnel_validation(ctx, param, value):
         return value
 
 
+def lookup_instance(region, ip_address):
+    filters = [{"Name": "network-interface.addresses.private-ip-address",
+                "Values": [str(ip_address)]}]
+    return next(piu.utils.list_running_instances(region, filters), None)
+
+
 def _request_access(even_url, cacert, username, hostname, reason, remote_host,
                     lifetime, clip, connect, tunnel):
     data = {'username': username, 'hostname': hostname, 'reason': reason}
@@ -214,10 +219,12 @@ def cli(ctx, config_file):
 @click.option('--connect', help='Directly connect to the host', envvar='PIU_CONNECT', is_flag=True, default=False)
 @click.option('--tunnel', help='Tunnel to the host', envvar='PIU_TUNNEL',
               callback=tunnel_validation, metavar='LOCALPORT:REMOTEPORT')
+@click.option('--check/--no-check', help='Verify that the EC2 instance exists and wasn\'t shutdown',
+              envvar='PIU_CHECK_INSTANCE', default=True)
 @region_option
 @click.pass_obj
 def request_access(config_file, host, reason, reason_cont, even_url, odd_host, lifetime, interactive,
-                   insecure, clip, connect, tunnel, region):
+                   insecure, clip, connect, tunnel, region, check):
     '''Request SSH access to a single host'''
     config = load_config(config_file)
     even_url = even_url or config.get('even_url')
@@ -243,6 +250,12 @@ def request_access(config_file, host, reason, reason_cont, even_url, odd_host, l
 
     try:
         ip = ipaddress.ip_address(hostname)
+
+        if check and not interactive and ip in STUPS_CIDR:
+            instance = lookup_instance(region, ip)
+            if instance is None:
+                click.confirm("No running instances found for {}, do you still want to request access?".format(ip),
+                              abort=True)
     except ValueError:
         ip = None
 
@@ -302,42 +315,33 @@ def request_access_interactive(region, odd_host):
     region = click.prompt('AWS region', default=region)
     odd_host = click.prompt('Odd SSH bastion hostname', default=odd_host)
 
-    ec2 = boto3.resource('ec2', region_name=region)
-    reservations = ec2.instances.filter(
-                   Filters=[{'Name': 'instance-state-name', 'Values': ['running']}])
-    name = stack_name = stack_version = None
-    instance_list = []
-    for r in reservations:
-        name = stack_name = stack_version = None
-        if not r.tags:
-            continue
-        for tag in r.tags:
-            tag_key, tag_value = tag['Key'], tag['Value']
-            if tag_key == 'Name':
-                name = tag_value
-            elif tag_key == 'StackName':
-                stack_name = tag_value
-            elif tag_key == 'StackVersion':
-                stack_version = tag_value
-        if name and stack_name and stack_version:
-            instance_list.append({'name': name, 'stack_name': stack_name, 'stack_version': stack_version,
-                                  'instance_id': r.instance_id, 'private_ip': r.private_ip_address})
-    instance_count = len(instance_list)
+    all_instances = piu.utils.list_running_instances(region, [])
+
+    stack_instances = [instance for instance in all_instances
+                       if instance.name and instance.stack_name and instance.stack_version]
+
+    instance_count = len(stack_instances)
     if instance_count == 0:
         raise click.ClickException('No running instances were found.')
-    sorted_instance_list = sorted(instance_list, key=operator.itemgetter('stack_name', 'stack_version'))
-    {d.update({'index': idx}) for idx, d in enumerate(sorted_instance_list, start=1)}
+
+    stack_instances.sort(key=operator.attrgetter('stack_name', 'stack_version'))
+
     print()
-    print_table('index name stack_name stack_version private_ip instance_id'.split(), sorted_instance_list)
+    table_entries = [dict(index=idx, **instance._asdict()) for idx, instance in enumerate(stack_instances, start=1)]
+    print_table(
+        'index name stack_name stack_version private_ip instance_id'.split(),
+        table_entries)
     print()
+
     if instance_count > 1:
         allowed_choices = ["{}".format(n) for n in range(1, instance_count + 1)]
         instance_index = int(click.prompt('Choose an instance (1-{})'.format(instance_count),
                                           type=click.Choice(allowed_choices))) - 1
     else:
-        click.confirm('Connect to {}?'.format(sorted_instance_list[0]['name']), default=True, abort=True)
+        click.confirm('Connect to {}?'.format(stack_instances[0].name), default=True, abort=True)
         instance_index = 0
-    host = sorted_instance_list[instance_index]['private_ip']
+
+    host = stack_instances[instance_index].private_ip
     reason = click.prompt('Reason', default='Troubleshooting')
     return (host, odd_host, reason)
 

piu/utils.py

@@ -1,5 +1,8 @@
 import boto3
 import click
+from collections import namedtuple
+
+Instance = namedtuple('Instance', ['instance_id', 'name', 'stack_name', 'stack_version', 'private_ip'])
 
 
 def _hosted_zones(route53):
@@ -28,6 +31,26 @@ def find_odd_host(region):
                 return record['Name'].rstrip('.')
 
 
+def list_running_instances(region, filters):
+    """Generator that yields Instance records for running EC2 instances matching the
+       filter and region"""
+    ec2 = boto3.resource('ec2', region_name=region)
+    effective_filters = [{'Name': 'instance-state-name',
+                          'Values': ['running']}]
+    effective_filters.extend(filters)
+    for instance in ec2.instances.filter(Filters=effective_filters):
+        instance_id = instance.instance_id
+
+        raw_tags = instance.tags or {}
+        tags = {tag['Key']: tag['Value'] for tag in raw_tags}
+
+        yield Instance(instance_id,
+                       tags.get('Name'),
+                       tags.get('StackName'),
+                       tags.get('StackVersion'),
+                       instance.private_ip_address)
+
+
 def current_region():
     """Returns the current AWS region"""
     session = boto3.session.Session()

tests/test_cli.py

@@ -2,13 +2,22 @@
 from unittest.mock import MagicMock
 import zign.api
 from piu.cli import cli
+import piu.utils
 import pytest
 
 
+def mock_list_running_instances(monkeypatch, *instances):
+    def mock_fn(region, filter):
+        yield from instances
+
+    monkeypatch.setattr('piu.utils.list_running_instances', mock_fn)
+
+
 @pytest.fixture(autouse=True)
 def mock_aws(monkeypatch):
     monkeypatch.setattr('piu.utils.current_region', lambda: 'eu-central-1')
     monkeypatch.setattr('piu.utils.find_odd_host', lambda region: None)
+    mock_list_running_instances(monkeypatch)
     yield
 
 
@@ -35,15 +44,61 @@ def test_success(monkeypatch):
     monkeypatch.setattr('zign.api.get_token', MagicMock(return_value='123'))
     monkeypatch.setattr('requests.post', MagicMock(return_value=response))
 
-    result = expect_success(['myuser@127.31.0.1',
+    result = expect_success(['myuser@172.31.0.1',
                              '--lifetime=15',
                              '--even-url=https://localhost/',
                              '--odd-host=odd.example.org',
+                             '--no-check',
                              'my reason'],
                             catch_exceptions=False)
 
     assert response.text in result.output
 
+@pytest.mark.parametrize("address,instance_exists,input,succeeded", [
+    # Stups IP, instance found => success
+    ("172.31.0.11", True, '', True),
+
+    # Stups IP, no instance found, confirmed => success
+    ("172.31.0.11", False, 'y', True),
+
+    # Stups IP, no instance found, not confirmed => failure
+    ("172.31.0.11", False, 'n', False),
+
+    # Other IP => success
+    ("10.0.1.1", False, None, True),
+
+    # Hostname => success
+    ("foo.example.org", False, None, True),
+])
+def test_instance_check(monkeypatch, address, instance_exists, input, succeeded):
+    success_text = '**MAGIC-SUCCESS**'
+    request = MagicMock(return_value=MagicMock(status_code=200, text=success_text))
+    monkeypatch.setattr('zign.api.get_token', MagicMock(return_value='123'))
+    monkeypatch.setattr('requests.post', request)
+
+    if instance_exists:
+        mock_list_running_instances(
+            monkeypatch,
+            piu.utils.Instance('i-123456', 'stack1-0o1o0', 'stack2', '0o1o0', address))
+
+    result = CliRunner().invoke(cli,
+                                ['myuser@{}'.format(address),
+                                 '--lifetime=15',
+                                 '--even-url=https://localhost/',
+                                 '--odd-host=odd.example.org',
+                                  'my reason'],
+                                input=input,
+                                catch_exceptions=False)
+
+    if succeeded:
+        assert request.called
+        assert result.exit_code == 0
+        assert success_text in result.output
+    else:
+        assert not request.called
+        assert result.exit_code != 0
+        assert success_text not in result.output
+
 
 def test_bad_request(monkeypatch):
     response = MagicMock(status_code=400, text='**MAGIC-BAD-REQUEST**')
@@ -87,7 +142,7 @@ def test_dialog(monkeypatch):
     monkeypatch.setattr('requests.get', MagicMock(return_value=response))
     monkeypatch.setattr('socket.getaddrinfo', MagicMock())
 
-    result = expect_success(['--config-file=config.yaml', 'req', 'myuser@172.31.0.1',
+    result = expect_success(['--config-file=config.yaml', 'req', 'myuser@172.31.0.1', '--no-check',
                              'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
     assert response.text in result.output
 
@@ -100,7 +155,7 @@ def test_oauth_failure(monkeypatch):
     monkeypatch.setattr('socket.getaddrinfo', MagicMock())
     runner = CliRunner()
 
-    result = runner.invoke(cli, ['--config-file=config.yaml', 'req', 'myuser@172.31.0.1',
+    result = runner.invoke(cli, ['--config-file=config.yaml', 'req', 'myuser@172.31.0.1', '--no-check',
                                  'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
 
     assert result.exit_code == 500
@@ -199,20 +254,12 @@ def test_interactive_success(monkeypatch):
     ec2 = MagicMock()
     request_access = MagicMock(return_value=200)
 
-    response = []
-    response.append(MagicMock(**{'instance_id': 'i-123456',
-                                 'private_ip_address': '172.31.10.10',
-                                 'tags': [{'Key': 'Name', 'Value': 'stack1-0o1o0'},
-                                          {'Key': 'StackVersion', 'Value': '0o1o0'},
-                                          {'Key': 'StackName', 'Value': 'stack1'}]
-                                 }))
-    response.append(MagicMock(**{'instance_id': 'i-789012',
-                                 'private_ip_address': '172.31.10.20',
-                                 'tags': [{'Key': 'Name', 'Value': 'stack2-0o1o0'},
-                                          {'Key': 'StackVersion', 'Value': '0o2o0'},
-                                          {'Key': 'StackName', 'Value': 'stack2'}]
-                                 }))
-    ec2.instances.filter = MagicMock(return_value=response)
+
+    instances = [
+        piu.utils.Instance('i-123456', 'stack1-0o1o0', 'stack2', '0o1o0', '172.31.10.10'),
+        piu.utils.Instance('i-789012', 'stack1-0o1o0', 'stack2', '0o2o0', '172.31.10.20')]
+
+    mock_list_running_instances(monkeypatch, *instances)
     monkeypatch.setattr('boto3.resource', MagicMock(return_value=ec2))
     monkeypatch.setattr('piu.cli._request_access', request_access)
 
@@ -232,14 +279,8 @@ def test_interactive_single_instance_success(monkeypatch):
     ec2 = MagicMock()
     request_access = MagicMock(return_value=200)
 
-    response = []
-    response.append(MagicMock(**{'instance_id': 'i-123456',
-                                 'private_ip_address': '172.31.10.10',
-                                 'tags': [{'Key': 'Name', 'Value': 'stack1-0o1o0'},
-                                          {'Key': 'StackVersion', 'Value': '0o1o0'},
-                                          {'Key': 'StackName', 'Value': 'stack1'}]
-                                 }))
-    ec2.instances.filter = MagicMock(return_value=response)
+    instance = piu.utils.Instance('i-123456', 'stack1-0o1o0', 'stack1', '0o1o0', '172.31.10.10')
+    mock_list_running_instances(monkeypatch, instance)
     monkeypatch.setattr('boto3.resource', MagicMock(return_value=ec2))
     monkeypatch.setattr('piu.cli._request_access', request_access)