From 367ffb4a59c2ef9a19aec8af118bfe17785db7ef Mon Sep 17 00:00:00 2001 From: Justin Whaley Date: Sun, 14 Feb 2021 18:33:33 -0500 Subject: [PATCH] Added functionality for AWS SSM I noticed that the existing application had nearly everything needed to do SSM sessions. I have modified the application to allow the ability to use your existing profiles and connect via SSM. --- aws_fuzzy_finder/main.py | 77 ++++++++++++++++++++++-------------- aws_fuzzy_finder/settings.py | 2 + 2 files changed, 49 insertions(+), 30 deletions(-) diff --git a/aws_fuzzy_finder/main.py b/aws_fuzzy_finder/main.py index c20781a..ffb1436 100644 --- a/aws_fuzzy_finder/main.py +++ b/aws_fuzzy_finder/main.py @@ -14,10 +14,13 @@ ENV_USE_PUBLIC_DNS_OVER_IP, ENV_KEY_PATH, ENV_SSH_COMMAND_TEMPLATE, + ENV_SSM_COMMAND_TEMPLATE, + ENV_USE_SSM, ENV_SSH_USER, ENV_TUNNEL_SSH_USER, ENV_TUNNEL_KEY_PATH, AWS_REGIONS, + AWS_DEFAULT_PROFILE, SEPARATOR, LIBRARY_PATH, CACHE_DIR, @@ -36,7 +39,8 @@ @click.option('--tunnel/--no-tunnel', help="Tunnel to another machine") @click.option('--tunnel-key-path', default='~/.ssh/id_rsa', help="Path to your private key, default: ~/.ssh/id_rsa") @click.option('--tunnel-user', default='ec2-user', help="User to SSH with, default: ec2-user") -def entrypoint(use_private_ip, key_path, user, ip_only, no_cache, tunnel, tunnel_key_path, tunnel_user): +@click.option('--ssm', 'use_ssm', flag_value=True, help="Tell the tool internally find the instance id and use AWS SSM") +def entrypoint(use_private_ip, key_path, user, ip_only, no_cache, tunnel, tunnel_key_path, tunnel_user, use_ssm): if not os.path.exists(CACHE_DIR): os.makedirs(CACHE_DIR) @@ -64,7 +68,7 @@ def entrypoint(use_private_ip, key_path, user, ip_only, no_cache, tunnel, tunnel } cache.close() except Exception as e: - print('Exception occured while getting cache, getting instances from AWS api: %s' % e) + print('Exception occurred while getting cache, getting instances from AWS api: %s' % e) if cache: cache.close() boto_instance_data = {} @@ -85,38 +89,47 @@ def entrypoint(use_private_ip, key_path, user, ip_only, no_cache, tunnel, tunnel LIBRARY_PATH ) - username = ENV_SSH_USER or user or '' - if username: - username = '%s@' % (username) + chosen_host = choice(fuzzysearch_bash_command, use_ssm) - key = ENV_KEY_PATH or key_path or '' - if key: - key = '-i %s' % (key) - - chosen_host = choice(fuzzysearch_bash_command) - - if ip_only: - sys.stdout.write(chosen_host) - exit(0) - - ssh_command = ENV_SSH_COMMAND_TEMPLATE.format( - user=username, - key=key, - host=chosen_host, - ) - - if tunnel: - ssh_command += " -t " + ENV_SSH_COMMAND_TEMPLATE.format( - user=ENV_TUNNEL_SSH_USER or tunnel_user, - key=ENV_TUNNEL_KEY_PATH or tunnel_key_path, - host=choice(fuzzysearch_bash_command), + if use_ssm or ENV_USE_SSM: + ssm_command = ENV_SSM_COMMAND_TEMPLATE.format( + profile=AWS_DEFAULT_PROFILE, + target=chosen_host, + ) + print(ssm_command) + subprocess.call(ssm_command, shell=True, executable='/bin/bash') + else: + if ip_only: + sys.stdout.write(chosen_host) + exit(0) + + username = ENV_SSH_USER or user or '' + if username: + username = '%s@' % (username) + + key = ENV_KEY_PATH or key_path or '' + if key: + key = '-i %s' % (key) + + ssh_command = ENV_SSH_COMMAND_TEMPLATE.format( + user=username, + key=key, + host=chosen_host, ) - print(ssh_command) - subprocess.call(ssh_command, shell=True, executable='/bin/bash') + if tunnel: + ssh_command += " -t " + ENV_SSH_COMMAND_TEMPLATE.format( + user=ENV_TUNNEL_SSH_USER or tunnel_user, + key=ENV_TUNNEL_KEY_PATH or tunnel_key_path, + host=choice(fuzzysearch_bash_command), + ) + + print(ssh_command) + subprocess.call(ssh_command, shell=True, executable='/bin/bash') -def choice(fuzzysearch_bash_command): +def choice(fuzzysearch_bash_command, use_ssm): + output = "" # used to collect the value returned try: choice = subprocess.check_output( fuzzysearch_bash_command, @@ -126,7 +139,11 @@ def choice(fuzzysearch_bash_command): except subprocess.CalledProcessError: exit(1) - return choice.split(SEPARATOR)[1].rstrip() + if use_ssm: + output = choice.split(' ')[1].replace('(', '').replace(')', '').rstrip() + else: + output = choice.split(SEPARATOR)[1].rstrip() + return output if __name__ == '__main__': diff --git a/aws_fuzzy_finder/settings.py b/aws_fuzzy_finder/settings.py index 80ea054..44e094a 100644 --- a/aws_fuzzy_finder/settings.py +++ b/aws_fuzzy_finder/settings.py @@ -7,10 +7,12 @@ ENV_SSH_USER = os.getenv('AWS_FUZZ_USER') ENV_KEY_PATH = os.getenv('AWS_FUZZ_KEY_PATH') ENV_USE_PRIVATE_IP = os.getenv('AWS_FUZZ_PRIVATE_IP') +ENV_USE_SSM = os.getenv('AWS_FUZZ_SSM') # use AWS Secure Session Manager instead of ssh ENV_USE_PUBLIC_DNS_OVER_IP = os.getenv('AWS_FUZZ_DNS_OVER_IP', False) # use public DNS over IP (both public or private) ENV_TUNNEL_SSH_USER = os.getenv('AWS_FUZZ_TUNNEL_USER') ENV_TUNNEL_KEY_PATH = os.getenv('AWS_FUZZ_TUNNEL_KEY_PATH') ENV_SSH_COMMAND_TEMPLATE = os.getenv('AWS_FUZZ_SSH_COMMAND_TEMPLATE', "ssh {key} {user}{host}") +ENV_SSM_COMMAND_TEMPLATE = os.getenv('AWS_FUZZ_SSM_COMMAND_TEMPLATE', "aws ssm start-session --profile {profile} --target {target}") ENV_AWS_REGIONS = os.getenv('AWS_FUZZ_AWS_REGIONS', '') CACHE_EXPIRY_TIME = int(os.getenv('AWS_FUZZ_CACHE_EXPIRY', 3600)) CACHE_ENABLED = os.getenv('AWS_FUZZ_USE_CACHE', False)