Merge pull request #28 from avb1989/feature/tunnel

#27 Implementation for tunnel feature

piu/cli.py

@@ -18,6 +18,7 @@
 import time
 import yaml
 import zign.api
+import re
 
 from clickclick import error, AliasedGroup, print_table, OutputFormat
 from .error_handling import handle_exceptions
@@ -128,7 +129,15 @@ def print_version(ctx, param, value):
     ctx.exit()
 
 
-def _request_access(even_url, cacert, username, hostname, reason, remote_host, lifetime, user, password, clip, connect):
+def tunnel_validation(ctx, param, value):
+    if value and not re.match("^[0-9]{1,5}:[0-9]{1,5}$", value):
+        raise click.BadParameter('Tunnel needs to be in format localPort:remotePort')
+    else:
+        return value
+
+
+def _request_access(even_url, cacert, username, hostname, reason, remote_host,
+                    lifetime, user, password, clip, connect, tunnel):
     data = {'username': username, 'hostname': hostname, 'reason': reason}
     host_via = hostname
     if remote_host:
@@ -154,12 +163,18 @@ def _request_access(even_url, cacert, username, hostname, reason, remote_host, l
         ssh_command = ''
         if remote_host:
             ssh_command = 'ssh -o StrictHostKeyChecking=no {username}@{remote_host}'.format(**vars())
+            if tunnel:
+                ports = tunnel.split(':')
+                ssh_command = '-L {local_port}:{remote_host}:{remote_port}'.format(
+                    local_port=ports[0], remote_host=remote_host, remote_port=ports[1])
         command = 'ssh -tA {username}@{hostname} {ssh_command}'.format(
                   username=username, hostname=hostname, ssh_command=ssh_command)
-        if connect:
+        if connect or tunnel:
             subprocess.call(command.split())
-        click.secho('You can now access your server with the following command:')
+
+        click.secho('You can access your server with the following command:')
         click.secho(command)
+
         if clip:
             click.secho('\nOr just check your clipboard and run ctrl/command + v (requires package "xclip" on Linux)')
             if pyperclip is not None:
@@ -195,9 +210,11 @@ def cli(ctx, config_file):
 @click.option('--insecure', help='Do not verify SSL certificate', is_flag=True, default=False)
 @click.option('--clip', help='Copy SSH command into clipboard', is_flag=True, default=False)
 @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.pass_obj
 def request_access(obj, host, reason, reason_cont, user, password, even_url, odd_host, lifetime, interactive,
-                   insecure, clip, connect):
+                   insecure, clip, connect, tunnel):
     '''Request SSH access to a single host'''
 
     if interactive:
@@ -207,6 +224,9 @@ def request_access(obj, host, reason, reason_cont, user, password, even_url, odd
     if not reason:
         raise click.UsageError('Missing argument "reason".')
 
+    if connect and tunnel:
+        raise click.UsageError('Cannot specify both "connect" and "tunnel"')
+
     user = user or zign.api.get_config().get('user') or os.getenv('USER')
 
     parts = host.split('@')
@@ -274,7 +294,7 @@ def request_access(obj, host, reason, reason_cont, user, password, even_url, odd
         remote_host = None
 
     return_code = _request_access(even_url, cacert, username, first_host, reason, remote_host, lifetime,
-                                  user, password, clip, connect)
+                                  user, password, clip, connect, tunnel)
 
     if return_code != 200:
         sys.exit(return_code)
@@ -328,7 +348,7 @@ def request_access_interactive():
     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
+                                          type=click.Choice(allowed_choices))) - 1
     else:
         click.confirm('Connect to {}?'.format(sorted_instance_list[0]['name']), default=True, abort=True)
         instance_index = 0

tests/test_cli.py

@@ -1,7 +1,6 @@
 
 from click.testing import CliRunner
 from unittest.mock import MagicMock
-import yaml
 import zign.api
 from piu.cli import cli
 
@@ -23,7 +22,14 @@ def test_success(monkeypatch):
     runner = CliRunner()
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['[email protected]', '--lifetime=15', '--even-url=https://localhost/', '--odd-host=odd.example.org', '--password=foobar', 'my reason'], catch_exceptions=False)
+        result = runner.invoke(cli,
+                               ['[email protected]',
+                                '--lifetime=15',
+                                '--even-url=https://localhost/',
+                                '--odd-host=odd.example.org',
+                                '--password=foobar',
+                                'my reason'],
+                               catch_exceptions=False)
 
     assert response.text in result.output
 
@@ -36,7 +42,14 @@ def test_bad_request(monkeypatch):
     runner = CliRunner()
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['req', '--lifetime=15', '--even-url=https://localhost/', '--password=foobar', 'myuser@odd-host', 'my reason'], catch_exceptions=False)
+        result = runner.invoke(cli,
+                               ['req',
+                                '--lifetime=15',
+                                '--even-url=https://localhost/',
+                                '--password=foobar',
+                                'myuser@odd-host',
+                                'my reason'],
+                               catch_exceptions=False)
 
     assert response.text in result.output
     assert 'Server returned status 400:' in result.output
@@ -50,7 +63,13 @@ def test_auth_failure(monkeypatch):
     runner = CliRunner()
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['r', '--even-url=https://localhost/', '--password=invalid', 'myuser@odd-host', 'my reason'], catch_exceptions=False)
+        result = runner.invoke(cli,
+                               ['r',
+                                '--even-url=https://localhost/',
+                                '--password=invalid',
+                                'myuser@odd-host',
+                                'my reason'],
+                               catch_exceptions=False)
 
     assert response.text in result.output
     assert 'Server returned status 403:' in result.output
@@ -67,11 +86,13 @@ def test_dialog(monkeypatch):
     runner = CliRunner()
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['--config-file=config.yaml', 'req', '[email protected]', 'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
+        result = runner.invoke(cli, ['--config-file=config.yaml', 'req', '[email protected]',
+                                     'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
 
     assert result.exit_code == 0
     assert response.text in result.output
 
+
 def test_oauth_failure(monkeypatch):
     response = MagicMock(status_code=200, text='**MAGIC-SUCCESS**')
     monkeypatch.setattr('zign.api.get_named_token', MagicMock(side_effect=zign.api.ServerError('**MAGIC-FAIL**')))
@@ -83,11 +104,13 @@ def test_oauth_failure(monkeypatch):
     runner = CliRunner()
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['--config-file=config.yaml', 'req', '[email protected]', 'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
+        result = runner.invoke(cli, ['--config-file=config.yaml', 'req', '[email protected]',
+                                     'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n')
 
     assert result.exit_code == 500
     assert 'Server error: **MAGIC-FAIL**' in result.output
 
+
 def test_login_arg_user(monkeypatch, tmpdir):
     arg_user = 'arg_user'
     zign_user = 'zign_user'
@@ -98,7 +121,7 @@ def test_login_arg_user(monkeypatch, tmpdir):
     runner = CliRunner()
 
     def mock__request_access(even_url, cacert, username, first_host, reason,
-                           remote_host, lifetime, user, password, clip):
+                             remote_host, lifetime, user, password, clip):
         assert arg_user == username
 
     monkeypatch.setattr('zign.api.get_config', lambda: {'user': zign_user})
@@ -107,8 +130,7 @@ def mock__request_access(even_url, cacert, username, first_host, reason,
     monkeypatch.setattr('requests.get', lambda x, timeout: response)
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access', '-U', arg_user],
-                catch_exceptions=False)
+        runner.invoke(cli, ['request-access', '-U', arg_user], catch_exceptions=False)
 
 
 def test_login_zign_user(monkeypatch, tmpdir):
@@ -120,7 +142,7 @@ def test_login_zign_user(monkeypatch, tmpdir):
     runner = CliRunner()
 
     def mock__request_access(even_url, cacert, username, first_host, reason,
-                          remote_host, lifetime, user, password, clip):
+                             remote_host, lifetime, user, password, clip):
         assert zign_user == username
 
     monkeypatch.setattr('zign.api.get_config', lambda: {'user': zign_user})
@@ -129,7 +151,7 @@ def mock__request_access(even_url, cacert, username, first_host, reason,
     monkeypatch.setattr('requests.get', lambda x, timeout: response)
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access'], catch_exceptions=False)
+        runner.invoke(cli, ['request-access'], catch_exceptions=False)
 
 
 def test_login_env_user(monkeypatch, tmpdir):
@@ -140,7 +162,7 @@ def test_login_env_user(monkeypatch, tmpdir):
     runner = CliRunner()
 
     def mock__request_access(even_url, cacert, username, first_host, reason,
-                          remote_host, lifetime, user, password, clip):
+                             remote_host, lifetime, user, password, clip):
         assert env_user == username
 
     monkeypatch.setattr('zign.api.get_config', lambda: {'user': ''})
@@ -149,26 +171,41 @@ def mock__request_access(even_url, cacert, username, first_host, reason,
     monkeypatch.setattr('requests.get', lambda x, timeout: response)
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access'], catch_exceptions=False)
+        runner.invoke(cli, ['request-access'], catch_exceptions=False)
 
 
 def test_interactive_success(monkeypatch):
     ec2 = MagicMock()
     request_access = MagicMock()
 
     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'}]}))
+    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)
-    boto3 = MagicMock()
     monkeypatch.setattr('boto3.resource', MagicMock(return_value=ec2))
     monkeypatch.setattr('piu.cli._request_access', MagicMock(side_effect=request_access))
 
     runner = CliRunner()
     input_stream = '\n'.join(['eu-west-1', '1', 'Troubleshooting']) + '\n'
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access', '--interactive', '--even-url=https://localhost/', '--odd-host=odd.example.org'], input=input_stream, catch_exceptions=False)
+        runner.invoke(cli,
+                      ['request-access',
+                       '--interactive',
+                       '--even-url=https://localhost/',
+                       '--odd-host=odd.example.org'],
+                      input=input_stream,
+                      catch_exceptions=False)
 
     assert request_access.called
 
@@ -178,37 +215,109 @@ def test_interactive_single_instance_success(monkeypatch):
     request_access = MagicMock()
 
     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-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)
-    boto3 = MagicMock()
     monkeypatch.setattr('boto3.resource', MagicMock(return_value=ec2))
     monkeypatch.setattr('piu.cli._request_access', MagicMock(side_effect=request_access))
 
     runner = CliRunner()
     input_stream = '\n'.join(['eu-west-1', '', 'Troubleshooting']) + '\n'
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access', '--interactive', '--even-url=https://localhost/', '--odd-host=odd.example.org'], input=input_stream, catch_exceptions=False)
+        runner.invoke(cli,
+                      ['request-access',
+                       '--interactive',
+                       '--even-url=https://localhost/',
+                       '--odd-host=odd.example.org'],
+                      input=input_stream,
+                      catch_exceptions=False)
 
     assert request_access.called
 
 
-
 def test_interactive_no_instances_failure(monkeypatch):
     ec2 = MagicMock()
     request_access = MagicMock()
 
     response = []
     ec2.instances.filter = MagicMock(return_value=response)
-    boto3 = MagicMock()
     monkeypatch.setattr('boto3.resource', MagicMock(return_value=ec2))
     monkeypatch.setattr('piu.cli._request_access', MagicMock(side_effect=request_access))
 
     runner = CliRunner()
     input_stream = '\neu-west-1\n'
 
     with runner.isolated_filesystem():
-        result = runner.invoke(cli, ['request-access', '--interactive', '--even-url=https://localhost/', '--odd-host=odd.example.org'], input=input_stream, catch_exceptions=False)
+        result = runner.invoke(cli,
+                               ['request-access',
+                                '--interactive',
+                                '--even-url=https://localhost/',
+                                '--odd-host=odd.example.org'],
+                               input=input_stream,
+                               catch_exceptions=False)
 
     assert result.exception
     assert 'Error: No running instances were found.' in result.output
+
+
+def test_tunnel_either_connect_or_tunnel():
+    input_stream = '\neu-central-1\n'
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli,
+                               ['request-access',
+                                '--connect',
+                                '--tunnel',
+                                '[email protected]',
+                                'Testing'],
+                               input=input_stream,
+                               catch_exceptions=False)
+    assert result.exception
+    assert 'Cannot specify both "connect" and "tunnel"'
+
+
+def test_tunnel_should_have_correct_format():
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli, ['request-access', '--tunnel', 'a2345:234a',
+                                     '[email protected]', 'Testing'], catch_exceptions=False)
+    assert result.exception
+
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli, ['request-access', '--tunnel', '23434',
+                                     '[email protected]', 'Testing'], catch_exceptions=False)
+    assert result.exception
+
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli, ['request-access', '--tunnel', 'a2345:2343',
+                                     '[email protected]', 'Testing'], catch_exceptions=False)
+    assert result.exception
+
+
+def test_tunnel_success(monkeypatch):
+
+    response = MagicMock(status_code=200, text='**MAGIC-SUCCESS**')
+
+    monkeypatch.setattr('zign.api.get_named_token', MagicMock(return_value={'access_token': '123'}))
+    monkeypatch.setattr('requests.post', MagicMock(return_value=response))
+    monkeypatch.setattr('subprocess.call', MagicMock())
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli, ['request-access',
+                                     '--tunnel', '2380:2379',
+                                     '--even-url=https://localhost/',
+                                     '--odd-host=odd.example.org',
+                                     '[email protected]',
+                                     'Testing'],
+                               catch_exceptions=False)
+
+    assert response.text in result.output
+    assert '-L 2380:somehost.example.org:2379' in result.output