diff --git a/README.md b/README.md index 2c75a6e..06c3844 100644 --- a/README.md +++ b/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 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 diff --git a/ssm_tools/resolver.py b/ssm_tools/resolver.py index c88ecb3..89c1ca3 100644 --- a/ssm_tools/resolver.py +++ b/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): diff --git a/ssm_tools/ssm_session_cli.py b/ssm_tools/ssm_session_cli.py index 76ec1a3..cf8c5ad 100755 --- a/ssm_tools/ssm_session_cli.py +++ b/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: