From 4eb4f067a724eca10ad207a9cebbf4453865278e Mon Sep 17 00:00:00 2001
From: Mark McIntyre <mark@mocktech.net>
Date: Wed, 4 Dec 2019 16:03:49 -0500
Subject: [PATCH] fixed a few bugs; added file permission change to private key

---
 rotate-ssh-keys | 44 ++++++++++++++++++++++++++++++++------------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/rotate-ssh-keys b/rotate-ssh-keys
index 51adcac..43c4d8a 100755
--- a/rotate-ssh-keys
+++ b/rotate-ssh-keys
@@ -4,17 +4,30 @@ import argparse
 import logging
 import boto3
 import time
+import os
 
 from cryptography.hazmat.primitives import serialization as crypto_serialization
 from cryptography.hazmat.primitives.asymmetric import rsa 
 from cryptography.hazmat.backends import default_backend as crypto_default_backend
 
+
+# setting up logging for this script
+_LEVEL = logging.INFO
+_FORMAT = "%(asctime)-15s [%(levelname)-8s] : %(lineno)d : %(name)s.%(funcName)s : %(message)s"
+logging.basicConfig(format=_FORMAT, level=_LEVEL)
+log = logging.getLogger()
+
+# set the boto logging levels to WARNING
+logging.getLogger('botocore').setLevel(logging.WARNING)
+logging.getLogger('boto3').setLevel(logging.WARNING)
+
+
 def parse_args():
     """
     Parse the arguments passed
     """
-    argp = argparser.ArgumentParser()
-    argp.add_argument('--debug', help="Run in debug mode")
+    argp = argparse.ArgumentParser()
+    argp.add_argument('--debug', action='store_true', help="Run in debug mode")
 
     argp.add_argument(
             '-p', '--profile',
@@ -39,7 +52,7 @@ def parse_args():
             help="String to use for the new key name and searching for existing keys"
     )
 
-    return args.parse_args()
+    return argp.parse_args()
 
 
 def get_session(profile_name=None, role_arn=None, region_name='us-east-1'):
@@ -124,7 +137,7 @@ def generate_ssh_keypair(key_size=2048, public_exponent=65537):
     return (public_key, private_key)
 
 
-def get_existing_keypair(session, prefix=""):
+def get_existing_keypairs(session, prefix=""):
     """
     Get the existing keypairs with the optional filter
     for a specific prefix of the key name
@@ -178,7 +191,7 @@ def upload_key(session, key_name, public_key):
         fingerprint = response['KeyFingerprint']
         log.info(f"Key fingerprint: {fingerprint}")
     except Exception as error:
-        log.error("Failed to upload key: {error}")
+        log.error(f"Failed to upload key: {error}")
 
     return fingerprint    
 
@@ -189,13 +202,22 @@ def main():
     if args.debug:
         log.setLevel(logging.DEBUG)
 
+        # let's keep the boto logging level sane
+        logging.getLogger('botocore').setLevel(logging.WARNING)
+        logging.getLogger('boto3').setLevel(logging.WARNING)
+
     log.info("Beginnging to generate new SSH key")
 
-    session = get_session()
+    session = get_session(profile_name=args.profile)
 
     # create the new key pair in memory
-    public_key, private_key = generate_ssh_key(args.key_size)
+    public_key, private_key = generate_ssh_keypair(args.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"{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:
@@ -204,15 +226,13 @@ def main():
     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)
     
-    # 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}")
-    
     # upload the new keypair to AWS account
     fingerprint = upload_key(session, key_name, public_key)