Initial check in for rmq agents for user mgmt, driver script to call them

block_account.py

+#!/usr/bin/env python3
+import json
+import rc_util
+import argparse
+import signal
+
+parser = argparse.ArgumentParser()
+parser.add_argument("username", help="username that will be created")
+parser.add_argument(
+    "-v", "--verbose", action="store_true", help="verbose output"
+)
+parser.add_argument(
+    "-n", "--dry-run", action="store_true", help="enable dry run mode"
+)
+args = parser.parse_args()
+
+timeout = 60
+
+callback_queue = rc_rmq.bind_queue(exclusive=True)
+
+def timeout_handler(signum, frame):
+    print("Process timeout, there's some issue with agents")
+    rc_util.rc_rmq.stop_consume()
+
+
+def callback(channel, method, properties, body):
+    msg = json.loads(body)
+    username = msg["username"]
+
+    if msg["success"]:
+        print(f"Account for {username} has been blocked :.")
+    else:
+        print(f"There's some issue while blocking account for {username}")
+        errmsg = msg.get("errmsg", [])
+        for err in errmsg:
+            print(err)
+
+    rc_util.rc_rmq.stop_consume()
+    rc_util.rc_rmq.delete_queue()
+
+
+rc_util.block_account(
+    args.username,
+    queuename=queuename,
+    email=args.email,
+    full=args.full_name,
+    reason=args.reason,
+)
+print(f"Lock action for {args.username} requested.")
+
+# Set initial timeout timer
+signal.signal(signal.SIGALRM, timeout_handler)
+signal.setitimer(signal.ITIMER_REAL, timeout)
+
+print("Waiting for completion...")
+rc_util.consume(
+    callback_queue, routing_key=f"complete.{queuename}", callback=callback
+)

prod_rmq_agents/block_new_jobs.py

+#!/usr/bin/env python
+import os
+import json
+import rc_util
+from pathlib import Path
+from rc_rmq import RCRMQ
+import rabbit_config as rcfg
+
+task = "block_new_jobs"
+
+args = rc_util.get_args()
+logger = rc_util.get_logger(args)
+
+# Instantiate rabbitmq object
+rc_rmq = RCRMQ({"exchange": "", "exchange_type": "topic"})
+
+
+def block_new_jobs(ch, method, properties, body):
+    msg = json.loads(body)
+    username = msg["username"]
+    msg["task"] = task
+    msg["success"] = False
+
+    reply_to = properties.reply_to
+
+    try:
+        block_ssh_cmd = f"/cm/shared/apps/slurm/19.05.5/bin/sacctmgr --immediate update user {username} set maxjobs=0"
+        block_ssh = popen(block_ssh_cmd).read().rstrip()
+
+    msg["success"] = True
+
+    except Exception:
+        msg["success"] = False
+        msg["errmsg"] = "Exception raised while setting maxjobs that can enter run state, check the logs for stack trace"
+        logger.error("", exc_info=True)
+
+    # send confirm message
+    rc_rmq.publish_msg(
+        {"routing_key": reply_to, 
+         "props": pika.BasicProperties(
+                            reply_to = callback_queue,
+                            ),
+         "msg": msg}
+    )
+
+    logger.debug(f"User {username} confirmation sent for {task}")
+
+    ch.basic_ack(delivery_tag=method.delivery_tag)
+
+
+logger.info(f"Start listening to queue: {task}")
+rc_rmq.start_consume(
+    {"queue": task, "routing_key": "block.*", "cb": block_new_jobs}
+)
+rc_rmq.start_consume(
+    {"queue": task, "routing_key": "block.newjobs.*", "cb": block_new_jobs}
+)
+
+logger.info("Disconnected")
+rc_rmq.disconnect()

prod_rmq_agents/lock_ssh.py

+#!/usr/bin/env python
+import os
+import json
+import rc_util
+from pathlib import Path
+from rc_rmq import RCRMQ
+import rabbit_config as rcfg
+
+task = "block_ssh"
+
+args = rc_util.get_args()
+logger = rc_util.get_logger(args)
+
+# Instantiate rabbitmq object
+rc_rmq = RCRMQ({"exchange": "", "exchange_type": "topic"})
+
+
+def block_ssh(ch, method, properties, body):
+    msg = json.loads(body)
+    username = msg["username"]
+    msg["task"] = task
+    msg["success"] = False
+
+    reply_to = properties.reply_to
+
+    block_ssh_cmd = f"/cm/local/apps/cmd/bin/cmsh -n -c "group; use nossh; append members {username}; commit;""
+    block_ssh = popen(block_ssh_cmd).read().rstrip()
+
+    msg["success"] = True
+
+    except Exception:
+        msg["success"] = False
+        msg["errmsg"] = "Exception raised, check the logs for stack trace"
+        logger.error("", exc_info=True)
+
+    # send confirm message
+    rc_rmq.publish_msg(
+        {"routing_key": reply_to,
+         "props": pika.BasicProperties(
+                            reply_to = callback_queue,
+                            ),
+         "msg": msg}
+    )
+
+    logger.debug(f"User {username} confirmation sent for {task}")
+
+    ch.basic_ack(delivery_tag=method.delivery_tag)
+
+
+logger.info(f"Start listening to queue: {task}")
+rc_rmq.start_consume(
+    {"queue": task, "routing_key": "block.*", "cb": block_ssh}
+)
+rc_rmq.start_consume(
+    {"queue": task, "routing_key": "block.ssh.*", "cb": block_ssh}
+)
+
+logger.info("Disconnected")
+rc_rmq.disconnect()