Fix key rotation bugs, support multi-user rotation, and improve logging/notifications

README.md

@@ -30,6 +30,12 @@ The config.yaml file need to be modified before used.
 $ python3 -m main.py -c "/path/to/config.yaml" -a rotate -u all
 ```
 
+`-u`/`--users` also accepts a single user or a comma-separated list:
+
+```bash
+$ python3 -m main.py -c "/path/to/config.yaml" -a rotate -u nagios,skydrop-azure
+```
+
 ## Running LOCK using Docker 🐳
 
 Pull the docker container:

project/main.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 LOCK_root = str(Path(__file__).resolve().parent.parent)
-print(f"Project is running from: [{LOCK_root}]")
+print(f"Project is running from: {LOCK_root}")
 
 import sys
 
@@ -25,10 +25,16 @@
 import requests
 
 
-def validate_keys_for_user(userdata, config_map, username, keys_to_delete):
+def resolve_target_users(all_users, usernames):
+    if usernames == "all":
+        return all_users
+    if isinstance(usernames, str):
+        usernames = [usernames]
+    return [u for u in all_users if next(iter(u)) in usernames]
+
+
+def validate_keys_for_user(userdata, config_map):
     username_to_validate = next(iter(userdata))
-    if username != "all" and username != username_to_validate:
-        return
     user_data = userdata.get(username_to_validate)
     if user_data.get("plugins"):
         iam_plugin = user_data.get("plugins")[0].get("iam")
@@ -42,7 +48,9 @@ def validate_keys_for_user(userdata, config_map, username, keys_to_delete):
                 )
                 if validation_result is not None:
                     old_key, prompt = validation_result
-                    keys_to_delete.append((username_to_validate, old_key, prompt))
+                    delete_old_key(
+                        user_data, config_map, username_to_validate, old_key, prompt
+                    )
             else:
                 logging.info(
                     f"   No get_new_key or rotate_ses_smtp_user section for iam plugin for user {username_to_validate} - skipping"
@@ -57,24 +65,15 @@ def validate_keys_for_user(userdata, config_map, username, keys_to_delete):
         )
 
 
-def validate_keys(username, all_users, config_map):
-    keys_to_delete = []
-    utils.run_threads(
-        all_users, validate_keys_for_user, config_map, username, keys_to_delete
-    )
-    for owner, key, prompt in keys_to_delete:
-        user_data = [data for data in all_users if next(iter(data)) == owner][0][owner]
-        delete_old_key(user_data, config_map, owner, key, prompt)
+def validate_keys(usernames, all_users, config_map):
+    target_users = resolve_target_users(all_users, usernames)
+    for user_data in target_users:
+        validate_keys_for_user(user_data, config_map)
 
 
-def rotate_update(
-    user_data, config_map, username=None, ssh_username=None, ssh_password=None
-):
-    if username is None:
-        username = next(iter(user_data))
-        modules = user_data[username]["plugins"]
-    else:
-        modules = user_data["plugins"]
+def rotate_update(user_data, config_map, ssh_username=None, ssh_password=None):
+    username = next(iter(user_data))
+    modules = user_data[username]["plugins"]
 
     update_access_key(username, ("", ""))
 
@@ -105,25 +104,21 @@ def rotate_update(
                 return
 
 
-def rotate_keys(username, all_users, config_map, user_data, ssh_username, ssh_password):
-    if username == "all":
-        utils.run_threads(
-            all_users, rotate_update, config_map, None, ssh_username, ssh_password
-        )
-    else:
-        rotate_update(user_data, config_map, username, ssh_username, ssh_password)
+def rotate_keys(usernames, all_users, config_map, ssh_username, ssh_password):
+    target_users = resolve_target_users(all_users, usernames)
+    utils.run_threads(
+        target_users, rotate_update, config_map, ssh_username, ssh_password
+    )
 
 
 def list_keys_for_user(user_data, config_map):
     username = next(iter(user_data))
     iam.list_keys(config_map, username)
 
 
-def list_keys(username, all_users, config_map):
-    if username == "all":
-        utils.run_threads(all_users, list_keys_for_user, config_map)
-    else:
-        iam.list_keys(config_map, username)
+def list_keys(usernames, all_users, config_map):
+    target_users = resolve_target_users(all_users, usernames)
+    utils.run_threads(target_users, list_keys_for_user, config_map)
 
 
 def update_access_key(username, key):
@@ -210,7 +205,12 @@ def main():
     parser = argparse.ArgumentParser(
         description="LOCK Let's Occasionally Circulate Keys"
     )
-    parser.add_argument("-u", "--user", help="aws user to rotate", required=False)
+    parser.add_argument(
+        "-u",
+        "--users",
+        help="aws user to rotate, or a comma-separated list of users",
+        required=False,
+    )
     parser.add_argument(
         "-c", "--config", help="Full path to a config file", required=True
     )
@@ -303,9 +303,12 @@ def main():
     verify_public_ip(public_ip_required)
 
     # args.dryRun = True
-    username = args.user
-    if args.user is None:
-        username = "test_lock"  # args.user
+    if args.users:
+        usernames = [u.strip() for u in args.users.split(",") if u.strip()]
+        if usernames == ["all"]:
+            usernames = "all"
+    else:
+        usernames = "test_lock"
     if args.action is None:
         args.action = "list"  # 'instance:status'
 
@@ -324,30 +327,32 @@ def main():
             ssh_password = input(f"Password for {args.ssh_username}: ")
 
     logging.debug(f"Config file {str(config_map)}")
-    user_data = None
     all_users = config_map["Users"]
-    for userdata in all_users:
-        if username == (next(iter(userdata))):
-            user_data = userdata.get(username)
 
-    if "user_data" not in locals() and username != "all":
-        logging.info(username + " does not exist in the config file.")
-        sys.exit()
+    if usernames != "all":
+        requested = usernames if isinstance(usernames, list) else [usernames]
+        all_usernames = {next(iter(u)) for u in all_users}
+        missing = [u for u in requested if u not in all_usernames]
+        if missing:
+            logging.info(f"{', '.join(missing)} does not exist in the config file.")
+            sys.exit()
 
     # get manually entered key, if any
     if args.key is not None:
-        update_access_key(username, args.key)
+        if isinstance(usernames, list) and len(usernames) != 1:
+            logging.error("-k/--key can only be used with a single user.")
+            sys.exit(1)
+        key_username = usernames[0] if isinstance(usernames, list) else usernames
+        update_access_key(key_username, args.key)
 
     if args.action == "list":
-        list_keys(username, all_users, config_map)
+        list_keys(usernames, all_users, config_map)
     elif args.action == "rotate":  # run functions listed in the config file.
-        rotate_keys(
-            username, all_users, config_map, user_data, args.ssh_username, ssh_password
-        )
+        rotate_keys(usernames, all_users, config_map, args.ssh_username, ssh_password)
     elif (
         args.action == "validate"
     ):  # validate that new key is being used and delete the old unused key
-        validate_keys(username, all_users, config_map)
+        validate_keys(usernames, all_users, config_map)
 
 
 if __name__ == "__main__":

project/plugins/1password.py

@@ -97,7 +97,7 @@ def create_item(
 
 def get_item_id(client: Client, vault_id: str, item_title: str) -> str:
     item_id = None
-    items = asyncio.run(client.items.list_all(vault_id)).obj
+    items = asyncio.run(client.items.list(vault_id))
     for item in items:
         if item.title == item_title:
             item_id = item.id
@@ -107,7 +107,7 @@ def get_item_id(client: Client, vault_id: str, item_title: str) -> str:
 
 def get_vault_id(client: Client, vault_title: str) -> str:
     vault_id = None
-    vaults = asyncio.run(client.vaults.list_all()).obj
+    vaults = asyncio.run(client.vaults.list())
     for vault in vaults:
         if vault.title == vault_title:
             vault_id = vault.id

project/plugins/bitbucket.py

@@ -42,9 +42,15 @@ def __get_bitbucket_token(config_map, username, **kwargs):
         "client_id": bb_api_key,
         "client_secret": bb_api_secret,
     }
-    access_token = requests.post(token_url, data=data).json()
-    api_token = access_token["access_token"]
-    return api_token
+    response = requests.post(token_url, data=data)
+    access_token = response.json()
+    if "access_token" not in access_token:
+        logging.error(
+            f"User {username}: Error retrieving Bitbucket API token "
+            f"(HTTP {response.status_code}): {access_token}"
+        )
+        return None
+    return access_token["access_token"]
 
 
 def __get_variable(api_token, workspace, variable_uuid):

project/plugins/iam.py

@@ -17,6 +17,14 @@ def get_iam_session():
     return boto3.Session(profile_name=values.profile)
 
 
+def get_profile_label(key_args):
+    if key_args.get("credential_profile") is not None:
+        return key_args.get("credential_profile")
+    if values.profile is not None:
+        return values.profile
+    return "static credentials"
+
+
 def get_iam_client(config_map, **kwargs):
     if kwargs.get("credential_profile") is not None:
         profile_name = kwargs.get("credential_profile")
@@ -384,14 +392,13 @@ def rotate_ses_smtp_user(config_map, username, **key_args):
 
             user_password = (key[0], password)
             update_user_password(user_password)
-            logging.info(f"User {username}: new user and password created")
             if values.hide_key is True:
-                print(
-                    f"                           New Username: {str(user_password[0])}"
+                logging.info(
+                    f"User {username}: New user and password created. New Username: {user_password[0]}"
                 )
             else:
-                print(
-                    f"                           New Username, Password: {str(user_password)}"
+                logging.info(
+                    f"User {username}: New user and password created. New Username, Password: {user_password}"
                 )
         else:
             logging.error(f"User {username}: Unable to get new key - skipping")
@@ -440,7 +447,8 @@ def store_password_parameter_store(config_map, username, **key_args):
             Overwrite=True,
         )
         logging.info(
-            f"User {username}: username and password written to parameter store."
+            f"User {username}: username and password written to parameter store "
+            f"in the '{get_profile_label(key_args)}' account."
         )
 
 
@@ -478,7 +486,8 @@ def store_key_parameter_store(config_map, username, **key_args):
             Overwrite=True,
         )
         logging.info(
-            f"User {username}: " + parameter_name + " key written to parameter store."
+            f"User {username}: {parameter_name} key written to parameter store "
+            f"in the '{get_profile_label(key_args)}' account."
         )
 
 
@@ -537,7 +546,6 @@ def get_ssm_client(config_map, **key_args):
 
     if key_args.get("credential_profile") is not None:
         profile_name = key_args.get("credential_profile")
-        print(profile_name)
         session = boto3.Session(profile_name=profile_name, region_name=region_name)
         return session.client("ssm")
     elif values.profile is not None:
@@ -557,7 +565,6 @@ def get_ecs_client(config_map, **key_args):
 
     if key_args.get("credential_profile") is not None:
         profile_name = key_args.get("credential_profile")
-        print(profile_name)
         session = boto3.Session(profile_name=profile_name, region_name=region_name)
         return session.client("ecs")
     elif values.profile is not None:

project/plugins/mail.py

@@ -7,6 +7,7 @@
 import boto3
 import logging
 import os
+import re
 
 
 def mail_message(config_map, username, **key_args):
@@ -50,13 +51,29 @@ def mail_message(config_map, username, **key_args):
         )
 
 
+URL_RE = re.compile(r"https?://\S+")
+
+
 class EmailTemplate:
     def __init__(self, template_name="", htmlvalues="", html=True, content_title=""):
         self.template_name = template_name
         self.htmlvalues = htmlvalues
         self.html = html
         self.content_title = content_title
 
+    @staticmethod
+    def _append_linkified(soup, parent, text):
+        last = 0
+        for match in URL_RE.finditer(text):
+            if match.start() > last:
+                parent.append(text[last : match.start()])
+            link = soup.new_tag("a", href=match.group(0))
+            link.string = match.group(0)
+            parent.append(link)
+            last = match.end()
+        if last < len(text):
+            parent.append(text[last:])
+
     def render(self):
         path = os.path.dirname(__file__)
         try:
@@ -69,7 +86,12 @@ def render(self):
             content1 = open(path + "/" + self.template_name).read()
 
         html = BeautifulSoup(content1, "html.parser")
-        html.find("div", {"id": "title"}).append(self.content_title)
+        title_div = html.find("div", {"id": "title"})
+        lines = re.split(r"\\n|\n", self.content_title)
+        for line in lines:
+            paragraph = html.new_tag("p", style="margin: 0 0 16px 0;")
+            self._append_linkified(html, paragraph, line.strip())
+            title_div.append(paragraph)
 
         return str(html)
 
@@ -124,7 +146,11 @@ def get_message(self):
 
 
 def send_ses(username, config_map, mail_msg):
-    if values.profile is not None:
+    credential_profile = config_map["Global"]["mail"].get("credential_profile")
+    if credential_profile is not None:
+        session = boto3.Session(profile_name=credential_profile, region_name="us-east-1")
+        ses = session.client("ses")
+    elif values.profile is not None:
         session = boto3.Session(profile_name=values.profile, region_name="us-east-1")
         ses = session.client("ses")
     else:

project/plugins/notification-email.html

@@ -27,7 +27,7 @@
                                 <tr>
                                     <td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
                                     <td><img
-                                            src="https://www.signiant.com/wp-content/themes/signiant-responsive/assets/img/logo-signiant@2x.png"
+                                            src="https://www.signiant.com/wp-content/uploads/2025/07/logo-1.svg"
                                             width="150" height="50"/></td>
                                 </tr>
                                 </tbody>
@@ -70,7 +70,7 @@
         <tr>
             <td class="content-block"
                 style="font-family: sans-serif; vertical-align: top; padding-bottom: 10px; padding-top: 10px; font-size: 12px; color: #999999; text-align: center;">
-                <p>Signiant,&nbsp;<span>11 Hines Rd, Kanata, ON K2K 1X7</span><br/><br/></p>
+                <p>Signiant,&nbsp;<span>4000 Innovation Dr, Suite 103, Kanata, ON K2K 3K1</span><br/><br/></p>
             </td>
         </tr>
         <tr>