Instance list cache

README.md

@@ -92,6 +92,48 @@ 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*.
+    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
+    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
+
+    ~ $ 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:
 
     This opens an interactive shell session over SSM without the need for
@@ -116,6 +158,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

requirements.txt

@@ -1,5 +1,6 @@
 pexpect
 packaging
+appdirs
 
 # AWS Stuff
 botocore

ssm_tools/resolver.py

@@ -7,6 +7,9 @@
 import botocore.credentials
 import botocore.session
 import boto3
+import appdirs
+import time
+import json
 
 logger = logging.getLogger("ssm-tools.resolver")
 
@@ -18,21 +21,84 @@ 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.use_cache = args.use_cache
+        self.cache_file = args.cache_file
+
+        if self.use_cache:
+            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
+
     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')
+
+
+        if self.use_cache:
+            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):
+            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:
+            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])
 
+        if self.use_cache:
+            items = self.get_list_cache()
+            if items is not None:
+                return items
+
         items = {}
 
         # List instances from SSM
@@ -105,6 +171,8 @@ def _try_append(_list, _dict, _key):
                                     items[instance_id]['InstanceName'] = tag['Value']
 
                             logger.debug("Updated instance: %s: %r", instance_id, items[instance_id])
+                    if self.use_cache:
+                        self.store_list_cache(items)
                     return items
 
             except botocore.exceptions.ClientError as ex:
@@ -123,6 +191,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.")
 
+        if self.use_cache:
+            self.store_list_cache(items)
         return items
 
     def print_list(self):

ssm_tools/ssm_session_cli.py

@@ -37,6 +37,9 @@ 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_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)')
@@ -62,8 +65,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 +114,19 @@ def main():
 
     try:
         instance = None
+
+        if args.update_cache:
+            args.use_cache = True
+            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: