From 57d7a46bdbb686dba16b6cf39acca0131d37b7a1 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Fri, 26 Aug 2022 14:31:11 +0300 Subject: [PATCH 1/8] store instances in json file in temp location --- ssm_tools/resolver.py | 50 ++++++++++++++++++++++++++++++++++++ ssm_tools/ssm_session_cli.py | 14 ++++++++-- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index f558aa4..d2446ca 100644 --- a/ssm_tools/resolver.py +++ b/ssm_tools/resolver.py @@ -7,6 +7,9 @@ import botocore.credentials import botocore.session import boto3 +import tempfile +import time +import json logger = logging.getLogger("ssm-tools.resolver") @@ -19,20 +22,65 @@ def __init__(self, args): self.session = boto3.session.Session(profile_name=args.profile, region_name=args.region) self.session._session.get_component('credential_provider').get_provider('assume-role').cache = botocore.credentials.JSONFileCache(cli_cache) + # Cache for ssm_tools + self.cache_dir = tempfile.gettempdir() class InstanceResolver(CommonResolver): + RESOLVER_CACHE_DURATION = 86400 # seconds + def __init__(self, args): super().__init__(args) # Create boto3 clients from session self.ssm_client = self.session.client('ssm') self.ec2_client = self.session.client('ec2') + + self.instance_cache = os.path.join(self.cache_dir, 'ssm_instances') + + def get_list_cache(self): + if not os.path.exists(self.instance_cache): + logger.debug('Did not load cache because cache file does not exist') + return None + + if time.time() - os.path.getmtime(self.instance_cache) > self.RESOLVER_CACHE_DURATION: + logger.debug('Did not load cache because file is older that one day') + return None + + with open(self.instance_cache, 'r') as fd: + try: + cache_content = json.load(fd) + except IOError: + logger.warning('Failed to load instance list cache due to IOError') + return None + + return cache_content + + def store_list_cache(self, ssm_list): + with open(self.instance_cache, 'w') as fd: + try: + json.dump(ssm_list, fd, indent=4) + except IOError: + logger.warning('Failed to update instance list cache due to IOError') + + def update_list_cache(self): + try: + os.remove(self.instance_cache) + except Exception as e: + print(e) + logger.warning('Failed to delete instance cache') + self.get_list() + + def get_list(self): def _try_append(_list, _dict, _key): if _key in _dict: _list.append(_dict[_key]) + items = self.get_list_cache() + if items is not None: + return items + items = {} # List instances from SSM @@ -102,6 +150,7 @@ def _try_append(_list, _dict, _key): items[instance_id]['InstanceName'] = tag['Value'] logger.debug("Updated instance: %s: %r", instance_id, items[instance_id]) + self.store_list_cache(items) return items except botocore.exceptions.ClientError as ex: @@ -120,6 +169,7 @@ def _try_append(_list, _dict, _key): if not tries: logger.warning("Unable to list instance details. Some instance names and IPs may be missing.") + self.store_list_cache(items) return items def print_list(self): diff --git a/ssm_tools/ssm_session_cli.py b/ssm_tools/ssm_session_cli.py index 76ec1a3..1967d07 100755 --- a/ssm_tools/ssm_session_cli.py +++ b/ssm_tools/ssm_session_cli.py @@ -37,6 +37,7 @@ def parse_args(argv): group_instance = parser.add_argument_group('Instance Selection') group_instance.add_argument('INSTANCE', nargs='?', help='Instance ID, Name, Host name or IP address') group_instance.add_argument('--list', '-l', dest='list', action="store_true", help='List instances available for SSM Session') + group_instance.add_argument('--update-cache', dest='update_cache', action="store_true", help='Updates local cache of instances available for SSM Session') group_session = parser.add_argument_group('Session Parameters') group_session.add_argument('--user', '-u', '--sudo', dest='user', metavar="USER", help='SUDO to USER after opening the session. Can\'t be used together with --document-name / --parameters. (optional)') @@ -62,8 +63,9 @@ def parse_args(argv): if args.show_version: show_version(args) - # Require exactly one of INSTANCE or --list - if bool(args.INSTANCE) + bool(args.list) != 1: + # Require exactly one of INSTANCE or --list and allow to do only cache update + if bool(args.INSTANCE) and bool(args.list) or \ + not args.update_cache and not bool(args.INSTANCE) and not bool(args.list): parser.error("Specify either INSTANCE or --list") if args.parameters and not args.document_name: @@ -110,10 +112,18 @@ def main(): try: instance = None + + if args.update_cache: + InstanceResolver(args).update_list_cache() + if args.list: InstanceResolver(args).print_list() sys.exit(0) + # If user wants to only update cache, no instance is specified and we should exit + if args.INSTANCE is None: + sys.exit(0) + instance_id = InstanceResolver(args).resolve_instance(args.INSTANCE) if not instance_id: From 4c9650165da5a775e556cd555d9c505f3ba6a621 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Fri, 26 Aug 2022 15:12:46 +0300 Subject: [PATCH 2/8] update README --- README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/README.md b/README.md index 2c75a6e..c3d043e 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,19 @@ and for ECS Docker Exec: `ecs-session` i-0beb42b1e6b60ac10 uswest2.aws.nz uswest2 172.31.0.92 ``` + If you don't see your instance in the list, there is a possibility + that is not stored in the local cache. You can update the local cache + with the flag `--update-cache`. *The cache expires after 1 day*. + + ``` + ~ $ ssm-session --list + i-07c189021bc56e042 test1.aws.nz test1 192.168.45.158 + + ~ $ ssm-session --list --update-cache + i-07c189021bc56e042 test1.aws.nz test1 192.168.45.158 + i-094df06d3633f3267 tunnel-test.aws.nz tunnel-test 192.168.44.95 + ``` + 2. **Open SSM session** to an instance: This opens an interactive shell session over SSM without the need for From 1fc498d6ead4133819f596448c9ebf4f90f08bfa Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Fri, 26 Aug 2022 15:30:48 +0300 Subject: [PATCH 3/8] remove forgotten print statement --- ssm_tools/resolver.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index d2446ca..38f67de 100644 --- a/ssm_tools/resolver.py +++ b/ssm_tools/resolver.py @@ -65,8 +65,7 @@ def store_list_cache(self, ssm_list): def update_list_cache(self): try: os.remove(self.instance_cache) - except Exception as e: - print(e) + except: logger.warning('Failed to delete instance cache') self.get_list() From ac606ae453ac182224b8f514c200c5bb09b269c2 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Mon, 26 Sep 2022 15:03:41 +0300 Subject: [PATCH 4/8] add appdirs in requirements --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 0444993..3e8b60d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ pexpect packaging +appdirs # AWS Stuff botocore From 8e020cf9507b0212bf3cb9512f4dfe1affef8f46 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Mon, 26 Sep 2022 15:04:00 +0300 Subject: [PATCH 5/8] change how and where instance cache is stored --- ssm_tools/resolver.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index 7063c25..ac1c425 100644 --- a/ssm_tools/resolver.py +++ b/ssm_tools/resolver.py @@ -7,7 +7,7 @@ import botocore.credentials import botocore.session import boto3 -import tempfile +import appdirs import time import json @@ -21,9 +21,14 @@ def __init__(self, args): # Construct boto3 session with MFA cache self.session = boto3.session.Session(profile_name=args.profile, region_name=args.region) self.session._session.get_component('credential_provider').get_provider('assume-role').cache = botocore.credentials.JSONFileCache(cli_cache) + + self.account_region = self.session.region_name + self.account_id = self.session.resource('iam').CurrentUser().user_id # Cache for ssm_tools - self.cache_dir = tempfile.gettempdir() + self.cache_dir = appdirs.user_cache_dir(appname="aws-ssm-tools") + if not os.path.isdir(self.cache_dir): + os.mkdir(self.cache_dir) class InstanceResolver(CommonResolver): RESOLVER_CACHE_DURATION = 86400 # seconds @@ -35,7 +40,7 @@ def __init__(self, args): self.ssm_client = self.session.client('ssm') self.ec2_client = self.session.client('ec2') - self.instance_cache = os.path.join(self.cache_dir, 'ssm_instances') + self.instance_cache = os.path.join(self.cache_dir, f'ssm_instances_{self.account_id}_{self.account_region}') def get_list_cache(self): if not os.path.exists(self.instance_cache): From 4d814499cb9d6606563d98b3f0df9544bf5d654c Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Mon, 26 Sep 2022 15:14:12 +0300 Subject: [PATCH 6/8] make instance cache use explicit --- ssm_tools/resolver.py | 24 +++++++++++++++--------- ssm_tools/ssm_session_cli.py | 3 ++- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index ac1c425..2b5c7e4 100644 --- a/ssm_tools/resolver.py +++ b/ssm_tools/resolver.py @@ -26,9 +26,11 @@ def __init__(self, args): self.account_id = self.session.resource('iam').CurrentUser().user_id # Cache for ssm_tools - self.cache_dir = appdirs.user_cache_dir(appname="aws-ssm-tools") - if not os.path.isdir(self.cache_dir): - os.mkdir(self.cache_dir) + self.use_cache = args.use_cache + if self.use_cache: + self.cache_dir = appdirs.user_cache_dir(appname="aws-ssm-tools") + if not os.path.isdir(self.cache_dir): + os.mkdir(self.cache_dir) class InstanceResolver(CommonResolver): RESOLVER_CACHE_DURATION = 86400 # seconds @@ -40,7 +42,8 @@ def __init__(self, args): self.ssm_client = self.session.client('ssm') self.ec2_client = self.session.client('ec2') - self.instance_cache = os.path.join(self.cache_dir, f'ssm_instances_{self.account_id}_{self.account_region}') + if self.use_cache: + self.instance_cache = os.path.join(self.cache_dir, f'ssm_instances_{self.account_id}_{self.account_region}') def get_list_cache(self): if not os.path.exists(self.instance_cache): @@ -81,9 +84,10 @@ def _try_append(_list, _dict, _key): if _key in _dict: _list.append(_dict[_key]) - items = self.get_list_cache() - if items is not None: - return items + if self.use_cache: + items = self.get_list_cache() + if items is not None: + return items items = {} @@ -157,7 +161,8 @@ def _try_append(_list, _dict, _key): items[instance_id]['InstanceName'] = tag['Value'] logger.debug("Updated instance: %s: %r", instance_id, items[instance_id]) - self.store_list_cache(items) + if self.use_cache: + self.store_list_cache(items) return items except botocore.exceptions.ClientError as ex: @@ -176,7 +181,8 @@ def _try_append(_list, _dict, _key): if not tries: logger.warning("Unable to list instance details. Some instance names and IPs may be missing.") - self.store_list_cache(items) + if self.use_cache: + self.store_list_cache(items) return items def print_list(self): diff --git a/ssm_tools/ssm_session_cli.py b/ssm_tools/ssm_session_cli.py index 1967d07..8e7fbea 100755 --- a/ssm_tools/ssm_session_cli.py +++ b/ssm_tools/ssm_session_cli.py @@ -44,7 +44,7 @@ def parse_args(argv): group_session.add_argument('--command', '-c', dest='command', metavar='COMMAND', help='Command to run in the SSM Session. Can\'t be used together with --user. If you need to run the COMMAND as a different USER prepend the command with the appropriate "sudo -u USER ...". (optional)') group_session.add_argument('--document-name', dest='document_name', help='Document to execute, e.g. AWS-StartInteractiveCommand (optional)') group_session.add_argument('--parameters', dest='parameters', help='Parameters for the --document-name, e.g. \'command=["sudo -i -u ec2-user"]\' (optional)') - + group_session.add_argument('--use-cache', dest='use_cache', action="store_true", help='Use locally stored cache of instances instead of hitting the AWS API') parser.description = 'Start SSM Shell Session to a given instance' parser.epilog = f''' IMPORTANT: instances must be registered in AWS Systems Manager (SSM) @@ -114,6 +114,7 @@ def main(): instance = None if args.update_cache: + args.use_cache = True InstanceResolver(args).update_list_cache() if args.list: From 8593575d7eb573d227ed6567d866497927377fc3 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Mon, 26 Sep 2022 15:27:03 +0300 Subject: [PATCH 7/8] update readme --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index c3d043e..51da2c0 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,9 @@ and for ECS Docker Exec: `ecs-session` i-0beb42b1e6b60ac10 uswest2.aws.nz uswest2 172.31.0.92 ``` + You can also store the fetched instances in your local cache to avoid + consecutive requests to AWS API. Cache is explicitly enabled with the + `--use-cache` flag. If you don't see your instance in the list, there is a possibility that is not stored in the local cache. You can update the local cache with the flag `--update-cache`. *The cache expires after 1 day*. @@ -129,6 +132,8 @@ and for ECS Docker Exec: `ecs-session` You can specify other SSM documents to run with `--document-name AWS-...` to customise your session. Refer to AWS docs for details. + You can also use your local cache to speed up instance resolving with `--use-cache`. + 3. **Open SSH session** over SSM with *port forwarding*. The `ssm-ssh` tool provides a connection and authentication mechanism From 06a55ed15077102903a23b598bedc3715c955150 Mon Sep 17 00:00:00 2001 From: Spyros Panagiotopoulos Date: Mon, 26 Sep 2022 15:45:37 +0300 Subject: [PATCH 8/8] add flag to use custom cache file --- README.md | 26 ++++++++++++++++++++++++++ ssm_tools/resolver.py | 18 ++++++++++++++---- ssm_tools/ssm_session_cli.py | 4 +++- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 51da2c0..06c3844 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,9 @@ and for ECS Docker Exec: `ecs-session` If you don't see your instance in the list, there is a possibility that is not stored in the local cache. You can update the local cache with the flag `--update-cache`. *The cache expires after 1 day*. + By default cache file is stored under the default user cache + directory of your operating system. You can specify your own cache + file with `--cache-file`. ``` ~ $ ssm-session --list @@ -106,6 +109,29 @@ and for ECS Docker Exec: `ecs-session` ~ $ ssm-session --list --update-cache i-07c189021bc56e042 test1.aws.nz test1 192.168.45.158 i-094df06d3633f3267 tunnel-test.aws.nz tunnel-test 192.168.44.95 + + ~ $ ssm-session --list --cache-file my_aws_cache + i-07c189021bc56e042 test1.aws.nz test1 192.168.45.158 + i-094df06d3633f3267 tunnel-test.aws.nz tunnel-test 192.168.44.95 + ~ $ cat my_aws_cache + { + "i-00cebf196bfa083ed": { + "InstanceId": "i-07c189021bc56e042", + "InstanceName": "test1", + "HostName": "test1.aws.nz", + "Addresses": [ + "192.168.45.158" + ] + }, + "i-00cebf196bfa083ed": { + "InstanceId": "i-094df06d3633f3267", + "InstanceName": "tunnel-test", + "HostName": "tunnel-test.aws.nz", + "Addresses": [ + "192.168.44.95" + ] + } + } ``` 2. **Open SSM session** to an instance: diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index 2b5c7e4..89c1ca3 100644 --- a/ssm_tools/resolver.py +++ b/ssm_tools/resolver.py @@ -27,10 +27,16 @@ def __init__(self, args): # Cache for ssm_tools self.use_cache = args.use_cache + self.cache_file = args.cache_file + if self.use_cache: - self.cache_dir = appdirs.user_cache_dir(appname="aws-ssm-tools") - if not os.path.isdir(self.cache_dir): - os.mkdir(self.cache_dir) + if self.cache_file is None: + self.cache_dir = appdirs.user_cache_dir(appname="aws-ssm-tools") + if not os.path.isdir(self.cache_dir): + os.mkdir(self.cache_dir) + else: + self.cache_dir = os.path.dirname(self.cache_file) + class InstanceResolver(CommonResolver): RESOLVER_CACHE_DURATION = 86400 # seconds @@ -42,8 +48,12 @@ def __init__(self, args): self.ssm_client = self.session.client('ssm') self.ec2_client = self.session.client('ec2') + if self.use_cache: - self.instance_cache = os.path.join(self.cache_dir, f'ssm_instances_{self.account_id}_{self.account_region}') + if self.cache_file is None: + self.instance_cache = os.path.join(self.cache_dir, f'ssm_instances_{self.account_id}_{self.account_region}') + else: + self.instance_cache = self.cache_file def get_list_cache(self): if not os.path.exists(self.instance_cache): diff --git a/ssm_tools/ssm_session_cli.py b/ssm_tools/ssm_session_cli.py index 8e7fbea..cf8c5ad 100755 --- a/ssm_tools/ssm_session_cli.py +++ b/ssm_tools/ssm_session_cli.py @@ -38,13 +38,15 @@ def parse_args(argv): group_instance.add_argument('INSTANCE', nargs='?', help='Instance ID, Name, Host name or IP address') group_instance.add_argument('--list', '-l', dest='list', action="store_true", help='List instances available for SSM Session') group_instance.add_argument('--update-cache', dest='update_cache', action="store_true", help='Updates local cache of instances available for SSM Session') + group_instance.add_argument('--use-cache', dest='use_cache', action="store_true", help='Use locally stored cache of instances instead of hitting the AWS API') + group_instance.add_argument('--cache-file', dest='cache_file', help='Custom path for the instance cache file') group_session = parser.add_argument_group('Session Parameters') group_session.add_argument('--user', '-u', '--sudo', dest='user', metavar="USER", help='SUDO to USER after opening the session. Can\'t be used together with --document-name / --parameters. (optional)') group_session.add_argument('--command', '-c', dest='command', metavar='COMMAND', help='Command to run in the SSM Session. Can\'t be used together with --user. If you need to run the COMMAND as a different USER prepend the command with the appropriate "sudo -u USER ...". (optional)') group_session.add_argument('--document-name', dest='document_name', help='Document to execute, e.g. AWS-StartInteractiveCommand (optional)') group_session.add_argument('--parameters', dest='parameters', help='Parameters for the --document-name, e.g. \'command=["sudo -i -u ec2-user"]\' (optional)') - group_session.add_argument('--use-cache', dest='use_cache', action="store_true", help='Use locally stored cache of instances instead of hitting the AWS API') + parser.description = 'Start SSM Shell Session to a given instance' parser.epilog = f''' IMPORTANT: instances must be registered in AWS Systems Manager (SSM)