runEccs.py

#!/usr/bin/env python3

import argparse
import asyncio
import datetime
import json
import time

import eccs_properties as e_p
import utils

from subprocess import PIPE

# Run Command
# https://docs.python.org/3/library/asyncio-queue.html#examples
# https://docs.python.org/3/library/asyncio-subprocess.html
async def run(name,queue,stdout_file,stderr_file,cmd_file):
   while True:
      # Get a "cmd item" out of the queue or wait for the next one
      cmd = await queue.get()

      # Elaborate "cmd" from shell.
      proc = await asyncio.create_subprocess_shell(
                   cmd,
                   stdout=asyncio.subprocess.PIPE,
                   stderr=asyncio.subprocess.PIPE
             )

      stdout, stderr = await proc.communicate()

      #print(f'[{name} exited with {proc.returncode}] - {cmd!r}')

      if stdout:
         stdout_file.write(f'-----\n[cmd]\n{cmd}\n\n[stdout]\n{stdout.decode()}')
      # If an error occurred, the failed command is put into the 'cmd_file' (failed-cmd.sh / failed-cmd-idp.sh) 
      if stderr:
         stderr_file.write(f'-----\n[cmd]\n{cmd}\n\n[stderr]\n{stderr.decode()}')
         cmd_file.write(f'{cmd}\n')

      # Notify the queue that the "work cmd" has been processed.
      queue.task_done()


async def main(cmd_list,stdout_file,stderr_file,cmd_file):
    # Create a queue that we will use to store our "workload".
    queue = asyncio.Queue()

    # Put all commands into the queue.
    for cmd in cmd_list:
        queue.put_nowait(cmd)

    # Create worker tasks to process the queue concurrently.
    tasks = []

    for i in range(e_p.ECCS_NUMPROCESSES):
        task = asyncio.create_task(run(f"cmd-{i}", queue, stdout_file, stderr_file, cmd_file))
        tasks.append(task)

    # Wait until the queue is fully processed.
    await queue.join()

    # Cancel our worker tasks.
    for task in tasks:
        task.cancel()

    # Wait until all worker tasks are cancelled.
    await asyncio.gather(*tasks, return_exceptions=True)


# MAIN
if __name__=="__main__":

   parser = argparse.ArgumentParser(description='This script will call another script in parallel to check one or more IdP on the correct consuming of the eduGAIN metadata')

   parser.add_argument("--idp", metavar="entityid", dest="idp_entityid", nargs=1, help="An IdP entityID")
   parser.add_argument("--test", action='store_true', dest="test", help="Test without effects")
   parser.add_argument("--replace", action='store_true', help="Check an IdP and replace the result")
      
   args = parser.parse_args()
         
   start = time.time()

   # Setup list_feds
   list_feds = utils.get_list_from_url(e_p.ECCS_LISTFEDSURL, e_p.ECCS_LISTFEDSFILE)

   # Setup list_eccs_idps
   list_eccs_idps = utils.get_list_from_url(e_p.ECCS_LISTIDPSURL, e_p.ECCS_LISTIDPSFILE)

   if (args.idp_entityid):
      stdout_file = open(e_p.ECCS_STDOUTIDP,"w+")
      stderr_file = open(e_p.ECCS_STDERRIDP,"w+")
      cmd_file = open(e_p.ECCS_FAILEDCMDIDP,"w+")
      idpJsonList = utils.get_idp_list(list_eccs_idps,idp_entityid=args.idp_entityid[0])

      if (args.test):
         cmd = f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idpJsonList[0])}' --test"
      elif (args.replace):
         cmd = f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idpJsonList[0])}' --replace"
      else:
         cmd = f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idpJsonList[0])}'"

      # List of only one command
      proc_list = [cmd]

      # Run Command
      asyncio.run(main(proc_list,stdout_file,stderr_file,cmd_file))

      # Close File
      stdout_file.close()
      stderr_file.close()
      cmd_file.close()

   else:
      stdout_file = open(e_p.ECCS_STDOUT,"w+")
      stderr_file = open(e_p.ECCS_STDERR,"w+")
      cmd_file = open(e_p.ECCS_FAILEDCMD,"w+")

      # Prepare input file for ECCS
      idpJsonList = utils.get_idp_list(list_eccs_idps)
      num_idps = len(idpJsonList)

      # Construct the list of commands to exec
      if (args.test):
         cmd_list = [[f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idp)}' --test"] for idp in idpJsonList]
      elif (args.replace):
         cmd_list = [[f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idp)}' --replace"] for idp in idpJsonList]
      else:
         cmd_list = [[f"{e_p.ECCS_DIR}/eccs.py '{json.dumps(idp)}'"] for idp in idpJsonList]

      # String Convertion needed for Asyncio
      proc_list = []
      count = 0
      while (count < num_idps):
         cmd = "".join(cmd_list.pop())
         proc_list.append(cmd)
         count = count + 1
 
      asyncio.run(main(proc_list,stdout_file,stderr_file,cmd_file))

      stdout_file.close()
      stderr_file.close()
      cmd_file.close()

   end = time.time()
   print("Time taken in hh:mm:ss - ", str(datetime.timedelta(seconds=end - start)))