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,36 +282,50 @@ 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}")
 
-    # create the new key pair in memory
+    else:
-    public_key, private_key = generate_ssh_keypair(args.key_size)
+        # handle the new key creation
+        log.info("Beginnging to generate new SSH key")
 
-    # get epoch of UTC time for the extension to make the name unique
+        # create the new key pair in memory
-    epoch_time = time.strftime("%s", time.gmtime())
+        public_key, private_key = generate_ssh_keypair(args.key_size)
-    key_name = f"{args.key_name_prefix}-{epoch_time}"
-    log.debug(f"key_name = {key_name}")
- 
-    # write the key values to files
-    log.info(f"Exporting the public key to {key_name}.pub")
-    with open(f"{key_name}.pub", 'w') as fp:
-        fp.write(public_key.decode('utf-8'))
 
-    log.info(f"Exporting the private key to file {key_name}")
+        # get epoch of UTC time for the extension to make the name unique
-    with open(key_name, 'w') as fp:
+        epoch_time = time.strftime("%s", time.gmtime())
-        fp.write(private_key.decode('utf-8'))
+        key_name = f"{args.key_name_prefix}-{epoch_time}"
+        log.debug(f"key_name = {key_name}")
+     
+        # write the key values to files
+        log.info(f"Exporting the public key to {key_name}.pub")
+        with open(f"{key_name}.pub", 'w') as fp:
+            fp.write(public_key.decode('utf-8'))
+
+        log.info(f"Exporting the private key to file {key_name}")
+        with open(key_name, 'w') as fp:
+            fp.write(private_key.decode('utf-8'))
+
+        log.debug("Setting permissions on private key file")
+        os.chmod(key_name, 0o600)
+        
+        # this list is for rotating the older keys out of circulation
+        existing_keypairs = get_existing_keypairs(session, args.key_name_prefix)
+        
+        # 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.debug("Setting permissions on private key file")
-    os.chmod(key_name, 0o600)
-    
-    # this list is for rotating the older keys out of circulation
-    existing_keypairs = get_existing_keypairs(session, args.key_name_prefix)
-    
-    # upload the new keypair to AWS account
-    fingerprint = upload_key(session, key_name, public_key)
-   
     log.info("Complete")