Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added functionality for AWS SSM #40

Merged
merged 1 commit into from
Feb 23, 2021
Merged

Added functionality for AWS SSM #40

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 47 additions & 30 deletions aws_fuzzy_finder/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Had to pull in the default profile from the environment variable so that you could change it needed.

SEPARATOR,
LIBRARY_PATH,
CACHE_DIR,
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed a typo there

cache.close()
boto_instance_data = {}
Expand All @@ -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:
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a conditional branch to just incorporate all the ssh vs ssm logic in one area

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,
Expand All @@ -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()
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

instead of stripping out the IP address, I have it pulling the instance id from the string.

else:
output = choice.split(SEPARATOR)[1].rstrip()
return output


if __name__ == '__main__':
Expand Down
2 changes: 2 additions & 0 deletions aws_fuzzy_finder/settings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just needed a different kind of template for find and replace.

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)
Expand Down