diff --git a/piu/cli.py b/piu/cli.py index 5fa0f8e..b231454 100644 --- a/piu/cli.py +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 8037df7..a5fa99f 100644 --- a/tests/test_cli.py +++ b/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, ['myuser@127.31.0.1', '--lifetime=15', '--even-url=https://localhost/', '--odd-host=odd.example.org', '--password=foobar', 'my reason'], catch_exceptions=False) + result = runner.invoke(cli, + ['myuser@127.31.0.1', + '--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', 'myuser@172.31.0.1', 'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n') + result = runner.invoke(cli, ['--config-file=config.yaml', 'req', 'myuser@172.31.0.1', + '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', 'myuser@172.31.0.1', 'my reason'], catch_exceptions=False, input='even\nodd\npassword\n\n') + result = runner.invoke(cli, ['--config-file=config.yaml', 'req', 'myuser@172.31.0.1', + '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,7 +171,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_interactive_success(monkeypatch): @@ -157,10 +179,19 @@ def test_interactive_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-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)) @@ -168,7 +199,13 @@ def test_interactive_success(monkeypatch): 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,9 +215,13 @@ 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)) @@ -188,19 +229,23 @@ def test_interactive_single_instance_success(monkeypatch): 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)) @@ -208,7 +253,71 @@ def test_interactive_no_instances_failure(monkeypatch): 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', + 'myuser@somehost.example.org', + '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', + 'myuser@somehost.example.org', 'Testing'], catch_exceptions=False) + assert result.exception + + with runner.isolated_filesystem(): + result = runner.invoke(cli, ['request-access', '--tunnel', '23434', + 'myuser@somehost.example.org', 'Testing'], catch_exceptions=False) + assert result.exception + + with runner.isolated_filesystem(): + result = runner.invoke(cli, ['request-access', '--tunnel', 'a2345:2343', + 'myuser@somehost.example.org', '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', + 'myuser@somehost.example.org', + 'Testing'], + catch_exceptions=False) + + assert response.text in result.output + assert '-L 2380:somehost.example.org:2379' in result.output