working on modifying the remote files

rotate-ssh-keys

@@ -10,6 +10,10 @@ 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.
+
+The epoch format of the time the keys are created make up the new
+key name. This all works when considering the latest key is the
+active key. Older keys will be removed when cleaning up keys.
 """
 
 
@@ -247,6 +251,47 @@ def upload_key(session, key_name, public_key):
     return fingerprint    
 
 
+def create_new_key(session, key_name_prefix, key_size):
+    """
+    Creates a new public and private key files and registers (uploads)
+    the new public key to AWS.
+
+    Args:
+        session: An AWS session object to establish access to the AWS account
+        key_name_prefix: Name of the key before the epoch value is appended
+        key_size: Size of the key to be created in bits
+    Returns:
+        Tuple of the key name and AWS key fingerprint
+    """
+    # 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(key_size)
+
+    # get epoch of UTC time for the extension to make the name unique
+    epoch_time = time.strftime("%s", time.gmtime())
+    key_name = f"{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)
+    
+    # upload the new keypair to AWS account
+    fingerprint = upload_key(session, key_name, public_key)
+
+    return (key_name, fingerprint)
+
+
 def remove_key(session, key_name):
     """
     Will remove the key to the AWS account with the provided name.
@@ -262,7 +307,7 @@ def remove_key(session, key_name):
     ec2_client = session.client('ec2')
 
     try:
-        log.info("Removing the key")
+        log.info(f"Removing the key {key_name}")
         response = ec2_client.delete_key_pair(KeyName=key_name)
         log.info(f"Key {key_name} successfully removed")
         return_value = True
@@ -272,6 +317,46 @@ def remove_key(session, key_name):
     return return_value
 
 
+def switch_keys_on_instances(session, new_key, old_key_name, remove_old_key_only=True):
+    """
+    Filters out the EC2 instances based on a specified old key names, adds the new
+    public key to the ~ec2-user/.ssh/authorized_keys file, and will remove the old
+    key in that file. If the remove_old_key_only is set to False, the
+    ~ec2-user/.ssh/authorized_keys file will be cleared with only the new present.
+
+    Args:
+        session: An AWS session object to establish access to the AWS account
+        new_key: String consisting of the public key with the key name appended
+        old_key_name: Key name (only) to filter out the EC2 instances to have the
+                new key added and all other keys removed
+        remove_old_key_only: Boolean to indicate whether the authorized_keys file
+                should only contain the new key or just remove the specified old
+                key (default is True)
+    Returns:
+        True if successful, False if unsuccessful
+    """
+    ec2_client = session.client('ec2')
+
+    instance_filters = {
+            Filters=[
+                {
+                    'Name': 'key-name',
+                    'Values': [
+                        old_key_name
+                    ]
+                }
+            ]
+    }
+
+    ec2_instances = [
+            y
+            for x in ec2_client.describe_instances(**instance_filters)['Reservations']
+            for y in x['Instances']
+    ]
+
+
+
+
 def main():
     args = parse_args()
 
@@ -286,45 +371,31 @@ def main():
     session = get_session(profile_name=args.profile, role_arn=args.role_arn)
 
     if ['delete-only', 'remove-only', 'order-66'] in args.action:
-        # handle specific key removal
+        # these actions are for key removal only
         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}")
+            raise RotateKeyException(f"The --removal_key_name argument must be provided with {args.action}")
 
     else:
-        # handle the new key creation
-        log.info("Beginnging to generate new SSH key")
+        # all other actions will create a new key and upload to AWS
+        key_name, fingerprint = create_new_key(session, args.key_name_prefix, args.key_size)
 
-        # create the new key pair in memory
-        public_key, private_key = generate_ssh_keypair(args.key_size)
+        # these actions will do a bit more by cleaning up
+        # the other keys after a new key is created
+        if ['lockdown', 'highlander', 'rotate'] in args.action:
+            log.info("Beginning key clean up phase")
+
+            # this list is for rotating the older keys out of circulation
+            existing_keypairs = get_existing_keypairs(session, args.key_name_prefix)
+            
+            for existing_keypair in existing_keypairs:
+                remove_key(session, existing_keypair)
 
-        # get epoch of UTC time for the extension to make the name unique
-        epoch_time = time.strftime("%s", time.gmtime())
-        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.info("Complete")