beginning to add the other components to this

rotate-ssh-keys

@@ -1,5 +1,18 @@
 #!/usr/bin/env python3
 
+"""
+This is designed to be run to manage SSH keys in AWS accounts 
+used for EC2 instances. It will handle key generation, AWS 
+registering (uploading), and rotating of the keys used on
+EC2 instances.
+
+The rotation process would on the creation of a new key identify
+all instances using the older key(s) and remove the public key
+values from the ~ec2-user/.ssh/authorized_keys file leaving only
+the public key value for the new key.
+"""
+
+
 import argparse
 import logging
 import boto3
@@ -21,20 +34,50 @@ log = logging.getLogger()
 logging.getLogger('botocore').setLevel(logging.WARNING)
 logging.getLogger('boto3').setLevel(logging.WARNING)
 
+# create a custom exception
+class RotateKeyError(Exception):
+    pass
+
 
 def parse_args():
     """
     Parse the arguments passed
     """
     argp = argparse.ArgumentParser()
-    argp.add_argument('--debug', action='store_true', help="Run in debug mode")
 
     argp.add_argument(
+            'action',
+            choices=[
+                'generate',
+                'add-only',
+                'highlander',
+                'lockdown',
+                'rotate',
+                'order-66',
+                'remove-only',
+                'delete-only',
+            ],
+            help="Action to perform: generate, add-only will create a new key and "
+                 "register with AWS. highlander, lockdown, rotate all will create "
+                 "a new key, register with AWS, and remove the old key from AWS "
+                 "and all the instances so that the key is rendered unusable. "
+                 "remove-only, delete-only, order-66 will remove the key specified "
+                 "by the --removal-key-name without generating a new key."
+    )
+
+    argp.add_argument(
+            '--debug',
+            action='store_true',
+            help="Run in debug mode"
+    )
+
+    account_access_group = argp.add_mutually_exclusive_group(required=True)
+    account_access_group.add_argument(
             '-p', '--profile',
             help="AWS profile name"
     )
 
-    argp.add_argument(
+    account_access_group.add_argument(
             '-r', '--role-arn',
             help="Role ARN with key access to the AWS account"
     )
@@ -52,6 +95,14 @@ def parse_args():
             help="String to use for the new key name and searching for existing keys"
     )
 
+    argp.add_argument(
+            '--removal-key-name',
+            default="",
+            help="Specific full name of the key to remove (must match exactly). "
+                 "This is ignored when action is not remove-only, delete-only, "
+                 "or order-66."
+    )
+
     return argp.parse_args()
 
 
@@ -196,6 +247,31 @@ def upload_key(session, key_name, public_key):
     return fingerprint    
 
 
+def remove_key(session, key_name):
+    """
+    Will remove the key to the AWS account with the provided name.
+
+    Args:
+        session: An AWS session object to establish access to the AWS account
+        key_name: Name for the keypair
+    Returns:
+        True if successful, False if unsuccessful
+    """
+    return_value = False
+    log.debug("Getting session client for EC2")
+    ec2_client = session.client('ec2')
+
+    try:
+        log.info("Removing the key")
+        response = ec2_client.delete_key_pair(KeyName=key_name)
+        log.info(f"Key {key_name} successfully removed")
+        return_value = True
+    except Exception as error:
+        log.error(f"Failed to upload key: {error}")
+
+    return return_value
+
+
 def main():
     args = parse_args()
 
@@ -206,9 +282,19 @@ def main():
         logging.getLogger('botocore').setLevel(logging.WARNING)
         logging.getLogger('boto3').setLevel(logging.WARNING)
 
-    log.info("Beginnging to generate new SSH key")
+    # get the session object to be used for all AWS access
+    session = get_session(profile_name=args.profile, role_arn=args.role_arn)
 
-    session = get_session(profile_name=args.profile)
+    if ['delete-only', 'remove-only', 'order-66'] in args.action:
+        # handle specific key removal
+        if args.removal_key_name):
+            remove_key(session, args.removal_key_name)
+        else:
+            raise RotateKeyException(f"--removal_key_name must be provided with {args.action}")
+
+    else:
+        # handle the new key creation
+        log.info("Beginnging to generate new SSH key")
 
         # create the new key pair in memory
         public_key, private_key = generate_ssh_keypair(args.key_size)
@@ -236,6 +322,10 @@ def main():
         # upload the new keypair to AWS account
         fingerprint = upload_key(session, key_name, public_key)
 
+    # handle the clean up of the other keys
+    if ['lockdown', 'highlander', 'rotate'] in args.action:
+        log.info("Beginning key clean up phase")
+
     log.info("Complete")