Add support for member SSH keys

code/workers/DH2AD/main.py

@@ -279,6 +279,31 @@ async def set_member_enabled(request: dict):
         logger.error(f"Error setting member enabled state: {e}")
         raise HTTPException(status_code=500, detail=f"Internal Server Error: " + str(e))
 
+@app.post("/v1/update_user_ssh_keys")
+async def update_user_ssh_keys(request: dict):
+    logger.debug(f"Ooo, gonna set user ssh keys: {request}")
+    try:
+        username = request.get("username")
+        ssh_keys = request.get("ssh_keys", [])
+        logger.info(f"Setting member ssh keys: {username}, ssh_keys={ssh_keys}")
+        result = perform_ad_operation("update_user_ssh_keys",
+                                      payload={"username", username, "ssh_keys",
+                                               ssh_keys})
+        if result is None:
+            raise HTTPException(status_code=500, detail="Failed to update member SSH keys in AD")
+        success, data = result
+        if data is None:
+            raise HTTPException(status_code=500, detail="Invalid response from Active Directory")
+        if not success:
+            logger.error(f"Failed to set member ssh keys: {data}")
+            raise HTTPException(status_code=500, detail="Failed to update SSH keys in AD")
+        if data.get("status") != "success":
+            logger.error(f"Failed to set member ssh keys: {data}")
+            raise HTTPException(status_code=500, detail="Failed to update SSH keys in AD")
+        return data["data"]
+    except Exception as e:
+        logger.error(f"Failed to set member ssh keys: {e}")
+        raise HTTPException(status_code=500, detail="Failed to update SSH keys in AD")
 
 ###############################################################################
 # Authorization management endpoints - these will call perform_ad_operation to add or
@@ -333,4 +358,4 @@ async def configure_authorizations(request: dict):
             }        
     except Exception as e:
         logger.error(f"Error configuring authorizations: {e}")
-        raise HTTPException(status_code=500, detail=f"Internal Server Error: " + str(e))
+        raise HTTPException(status_code=500, detail=f"Internal Server Error: " + str(e))

code/workers/DHADController/main.py

@@ -327,6 +327,34 @@ def remove_user_from_group(username, group_name):
             "error": str(e)
         }
 
+
+def update_user_ssh_keys(username, ssh_keys):
+    logger.debug(f"Updating user {username}'s SSH keys")
+    ad_session = ad.create_ad_session()    
+    try:
+        success = ad.update_user(ad_session, username,
+                                 {"altSecurityIdentities", ssh_keys})
+        if not success:
+            logger.error(f"Failed to update SSH keys for user: {username}")
+            return {
+                "status": "failure",
+                "error": f"Failed to update SSH keys for user: {username}",
+            }
+        logger.info(f"Updated SSH keys for user {username}")
+        return {
+            "status": "success",
+            "data": {
+                "operation": "update_user_ssh_keys",
+                "username": username,
+                "ssh_keys": ssh_keys
+            }
+        }
+    except Exception as e:
+        logger.error(f"Error updating user {username}'s ssh keys: {e}")
+        return {
+            "status": "failure",
+            "error": str(e)
+        }
 def sync_account_info(request):
     logger.debug(f"Syncing account information with request: {request}")
     # Okay, first let's see if there's already an AD user in Active Directory 
@@ -371,8 +399,7 @@ def sync_account_info(request):
             "status": "failure",
             "error": str(e)
         }
-
-
+
 ###############################################################################
 # Message Queue Interaction Functions
 ###############################################################################
@@ -435,6 +462,9 @@ def handle_message(msg_id, payload):
         return sync_account_info(payload)
     elif operation == "get_is_user_enabled":
         return get_is_user_enabled(username=payload.get("username"))
+    elif operation == "update_user_ssh_keys":
+        return update_user_ssh_keys(username=payload.get("username"),
+                                    ssh_keys=payload.get("ssh_keys"))
     else:
         result_data = {
             "original_id": msg_id,
@@ -518,4 +548,4 @@ def main():
     #result = create_user("zestyzest", "Zesty", "Zest", "zesty1234@example.com", "12345")
     #result = add_user_to_group("zestyzest", "Band Saw")
     #result = remove_user_from_group("zestyzest", "Band Saw")
-    #print(result)
+    #print(result)