1#!/usr/bin/env python 2 3r""" 4This module provides valuable argument processing functions like gen_get_options and sprint_args. 5""" 6 7import sys 8import os 9import re 10try: 11 import psutil 12 psutil_imported = True 13except ImportError: 14 psutil_imported = False 15try: 16 import __builtin__ 17except ImportError: 18 import builtins as __builtin__ 19import atexit 20import signal 21import argparse 22 23import gen_print as gp 24import gen_valid as gv 25import gen_cmd as gc 26import gen_misc as gm 27 28default_string = ' The default value is "%(default)s".' 29module = sys.modules["__main__"] 30 31 32def gen_get_options(parser, 33 stock_list=[]): 34 r""" 35 Parse the command line arguments using the parser object passed and return True/False (i.e. pass/fail). 36 However, if gv.exit_on_error is set, simply exit the program on failure. Also set the following built in 37 values: 38 39 __builtin__.quiet This value is used by the qprint functions. 40 __builtin__.test_mode This value is used by command processing functions. 41 __builtin__.debug This value is used by the dprint functions. 42 __builtin__.arg_obj This value is used by print_program_header, etc. 43 __builtin__.parser This value is used by print_program_header, etc. 44 45 Description of arguments: 46 parser A parser object. See argparse module documentation for details. 47 stock_list The caller can use this parameter to request certain stock parameters 48 offered by this function. For example, this function will define a 49 "quiet" option upon request. This includes stop help text and parm 50 checking. The stock_list is a list of tuples each of which consists of 51 an arg_name and a default value. Example: stock_list = [("test_mode", 52 0), ("quiet", 1), ("debug", 0)] 53 """ 54 55 # This is a list of stock parms that we support. 56 master_stock_list = ["quiet", "test_mode", "debug", "loglevel"] 57 58 # Process stock_list. 59 for ix in range(0, len(stock_list)): 60 if len(stock_list[ix]) < 1: 61 error_message = "Programmer error - stock_list[" + str(ix) +\ 62 "] is supposed to be a tuple containing at" +\ 63 " least one element which is the name of" +\ 64 " the desired stock parameter:\n" +\ 65 gp.sprint_var(stock_list) 66 return gv.process_error_message(error_message) 67 if isinstance(stock_list[ix], tuple): 68 arg_name = stock_list[ix][0] 69 default = stock_list[ix][1] 70 else: 71 arg_name = stock_list[ix] 72 default = None 73 74 if arg_name not in master_stock_list: 75 error_message = "Programmer error - arg_name \"" + arg_name +\ 76 "\" not found found in stock list:\n" +\ 77 gp.sprint_var(master_stock_list) 78 return gv.process_error_message(error_message) 79 80 if arg_name == "quiet": 81 if default is None: 82 default = 0 83 parser.add_argument( 84 '--quiet', 85 default=default, 86 type=int, 87 choices=[1, 0], 88 help='If this parameter is set to "1", %(prog)s' 89 + ' will print only essential information, i.e. it will' 90 + ' not echo parameters, echo commands, print the total' 91 + ' run time, etc.' + default_string) 92 elif arg_name == "test_mode": 93 if default is None: 94 default = 0 95 parser.add_argument( 96 '--test_mode', 97 default=default, 98 type=int, 99 choices=[1, 0], 100 help='This means that %(prog)s should go through all the' 101 + ' motions but not actually do anything substantial.' 102 + ' This is mainly to be used by the developer of' 103 + ' %(prog)s.' + default_string) 104 elif arg_name == "debug": 105 if default is None: 106 default = 0 107 parser.add_argument( 108 '--debug', 109 default=default, 110 type=int, 111 choices=[1, 0], 112 help='If this parameter is set to "1", %(prog)s will print' 113 + ' additional debug information. This is mainly to be' 114 + ' used by the developer of %(prog)s.' + default_string) 115 elif arg_name == "loglevel": 116 if default is None: 117 default = "info" 118 parser.add_argument( 119 '--loglevel', 120 default=default, 121 type=str, 122 choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL', 123 'debug', 'info', 'warning', 'error', 'critical'], 124 help='If this parameter is set to "1", %(prog)s will print' 125 + ' additional debug information. This is mainly to be' 126 + ' used by the developer of %(prog)s.' + default_string) 127 128 arg_obj = parser.parse_args() 129 130 __builtin__.quiet = 0 131 __builtin__.test_mode = 0 132 __builtin__.debug = 0 133 __builtin__.loglevel = 'WARNING' 134 for ix in range(0, len(stock_list)): 135 if isinstance(stock_list[ix], tuple): 136 arg_name = stock_list[ix][0] 137 default = stock_list[ix][1] 138 else: 139 arg_name = stock_list[ix] 140 default = None 141 if arg_name == "quiet": 142 __builtin__.quiet = arg_obj.quiet 143 elif arg_name == "test_mode": 144 __builtin__.test_mode = arg_obj.test_mode 145 elif arg_name == "debug": 146 __builtin__.debug = arg_obj.debug 147 elif arg_name == "loglevel": 148 __builtin__.loglevel = arg_obj.loglevel 149 150 __builtin__.arg_obj = arg_obj 151 __builtin__.parser = parser 152 153 # For each command line parameter, create a corresponding global variable and assign it the appropriate 154 # value. For example, if the command line contained "--last_name='Smith', we'll create a global variable 155 # named "last_name" with the value "Smith". 156 module = sys.modules['__main__'] 157 for key in arg_obj.__dict__: 158 setattr(module, key, getattr(__builtin__.arg_obj, key)) 159 160 return True 161 162 163def set_pgm_arg(var_value, 164 var_name=None): 165 r""" 166 Set the value of the arg_obj.__dict__ entry named in var_name with the var_value provided. Also, set 167 corresponding global variable. 168 169 Description of arguments: 170 var_value The value to set in the variable. 171 var_name The name of the variable to set. This defaults to the name of the 172 variable used for var_value when calling this function. 173 """ 174 175 if var_name is None: 176 var_name = gp.get_arg_name(None, 1, 2) 177 178 arg_obj.__dict__[var_name] = var_value 179 module = sys.modules['__main__'] 180 setattr(module, var_name, var_value) 181 if var_name == "quiet": 182 __builtin__.quiet = var_value 183 elif var_name == "debug": 184 __builtin__.debug = var_value 185 elif var_name == "test_mode": 186 __builtin__.test_mode = var_value 187 188 189def sprint_args(arg_obj, 190 indent=0): 191 r""" 192 sprint_var all of the arguments found in arg_obj and return the result as a string. 193 194 Description of arguments: 195 arg_obj An argument object such as is returned by the argparse parse_args() 196 method. 197 indent The number of spaces to indent each line of output. 198 """ 199 200 col1_width = gp.dft_col1_width + indent 201 202 buffer = "" 203 for key in arg_obj.__dict__: 204 buffer += gp.sprint_varx(key, getattr(arg_obj, key), 0, indent, 205 col1_width) 206 return buffer 207 208 209def sync_args(): 210 r""" 211 Synchronize the argument values to match their corresponding global variable values. 212 213 The user's validate_parms() function may manipulate global variables that correspond to program 214 arguments. After validate_parms() is called, sync_args is called to set the altered values back into the 215 arg_obj. This will ensure that the print-out of program arguments reflects the updated values. 216 217 Example: 218 219 def validate_parms(): 220 221 # Set a default value for dir_path argument. 222 dir_path = gm.add_trailing_slash(gm.dft(dir_path, os.getcwd())) 223 """ 224 module = sys.modules['__main__'] 225 for key in arg_obj.__dict__: 226 arg_obj.__dict__[key] = getattr(module, key) 227 228 229term_options = None 230 231 232def set_term_options(**kwargs): 233 r""" 234 Set the global term_options. 235 236 If the global term_options is not None, gen_exit_function() will call terminate_descendants(). 237 238 Description of arguments(): 239 kwargs Supported keyword options follow: 240 term_requests Requests to terminate specified descendants of this program. The 241 following values for term_requests are supported: 242 children Terminate the direct children of this program. 243 descendants Terminate all descendants of this program. 244 <dictionary> A dictionary with support for the following keys: 245 pgm_names A list of program names which will be used to identify which descendant 246 processes should be terminated. 247 """ 248 249 global term_options 250 # Validation: 251 arg_names = list(kwargs.keys()) 252 gv.valid_list(arg_names, ['term_requests']) 253 if type(kwargs['term_requests']) is dict: 254 keys = list(kwargs['term_requests'].keys()) 255 gv.valid_list(keys, ['pgm_names']) 256 else: 257 gv.valid_value(kwargs['term_requests'], ['children', 'descendants']) 258 term_options = kwargs 259 260 261if psutil_imported: 262 def match_process_by_pgm_name(process, pgm_name): 263 r""" 264 Return True or False to indicate whether the process matches the program name. 265 266 Description of argument(s): 267 process A psutil process object such as the one returned by psutil.Process(). 268 pgm_name The name of a program to look for in the cmdline field of the process 269 object. 270 """ 271 272 # This function will examine elements 0 and 1 of the cmdline field of the process object. The 273 # following examples will illustrate the reasons for this: 274 275 # Example 1: Suppose a process was started like this: 276 277 # shell_cmd('python_pgm_template --quiet=0', fork=1) 278 279 # And then this function is called as follows: 280 281 # match_process_by_pgm_name(process, "python_pgm_template") 282 283 # The process object might contain the following for its cmdline field: 284 285 # cmdline: 286 # [0]: /usr/bin/python 287 # [1]: /my_path/python_pgm_template 288 # [2]: --quiet=0 289 290 # Because "python_pgm_template" is a python program, the python interpreter (e.g. "/usr/bin/python") 291 # will appear in entry 0 of cmdline and the python_pgm_template will appear in entry 1 (with a 292 # qualifying dir path). 293 294 # Example 2: Suppose a process was started like this: 295 296 # shell_cmd('sleep 5', fork=1) 297 298 # And then this function is called as follows: 299 300 # match_process_by_pgm_name(process, "sleep") 301 302 # The process object might contain the following for its cmdline field: 303 304 # cmdline: 305 # [0]: sleep 306 # [1]: 5 307 308 # Because "sleep" is a compiled executable, it will appear in entry 0. 309 310 optional_dir_path_regex = "(.*/)?" 311 cmdline = process.as_dict()['cmdline'] 312 return re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[0]) \ 313 or re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[1]) 314 315 def select_processes_by_pgm_name(processes, pgm_name): 316 r""" 317 Select the processes that match pgm_name and return the result as a list of process objects. 318 319 Description of argument(s): 320 processes A list of psutil process objects such as the one returned by 321 psutil.Process(). 322 pgm_name The name of a program to look for in the cmdline field of each process 323 object. 324 """ 325 326 return [process for process in processes if match_process_by_pgm_name(process, pgm_name)] 327 328 def sprint_process_report(pids): 329 r""" 330 Create a process report for the given pids and return it as a string. 331 332 Description of argument(s): 333 pids A list of process IDs for processes to be included in the report. 334 """ 335 report = "\n" 336 cmd_buf = "echo ; ps wwo user,pgrp,pid,ppid,lstart,cmd --forest " + ' '.join(pids) 337 report += gp.sprint_issuing(cmd_buf) 338 rc, outbuf = gc.shell_cmd(cmd_buf, quiet=1) 339 report += outbuf + "\n" 340 341 return report 342 343 def get_descendant_info(process=psutil.Process()): 344 r""" 345 Get info about the descendants of the given process and return as a tuple of descendants, 346 descendant_pids and process_report. 347 348 descendants will be a list of process objects. descendant_pids will be a list of pids (in str form) 349 and process_report will be a report produced by a call to sprint_process_report(). 350 351 Description of argument(s): 352 process A psutil process object such as the one returned by psutil.Process(). 353 """ 354 descendants = process.children(recursive=True) 355 descendant_pids = [str(process.pid) for process in descendants] 356 if descendants: 357 process_report = sprint_process_report([str(process.pid)] + descendant_pids) 358 else: 359 process_report = "" 360 return descendants, descendant_pids, process_report 361 362 def terminate_descendants(): 363 r""" 364 Terminate descendants of the current process according to the requirements layed out in global 365 term_options variable. 366 367 Note: If term_options is not null, gen_exit_function() will automatically call this function. 368 369 When this function gets called, descendant processes may be running and may be printing to the same 370 stdout stream being used by this process. If this function writes directly to stdout, its output can 371 be interspersed with any output generated by descendant processes. This makes it very difficult to 372 interpret the output. In order solve this problem, the activity of this process will be logged to a 373 temporary file. After descendant processes have been terminated successfully, the temporary file 374 will be printed to stdout and then deleted. However, if this function should fail to complete (i.e. 375 get hung waiting for descendants to terminate gracefully), the temporary file will not be deleted and 376 can be used by the developer for debugging. If no descendant processes are found, this function will 377 return before creating the temporary file. 378 379 Note that a general principal being observed here is that each process is responsible for the 380 children it produces. 381 """ 382 383 message = "\n" + gp.sprint_dashes(width=120) \ 384 + gp.sprint_executing() + "\n" 385 386 current_process = psutil.Process() 387 388 descendants, descendant_pids, process_report = get_descendant_info(current_process) 389 if not descendants: 390 # If there are no descendants, then we have nothing to do. 391 return 392 393 terminate_descendants_temp_file_path = gm.create_temp_file_path() 394 gp.print_vars(terminate_descendants_temp_file_path) 395 396 message += gp.sprint_varx("pgm_name", gp.pgm_name) \ 397 + gp.sprint_vars(term_options) \ 398 + process_report 399 400 # Process the termination requests: 401 if term_options['term_requests'] == 'children': 402 term_processes = current_process.children(recursive=False) 403 term_pids = [str(process.pid) for process in term_processes] 404 elif term_options['term_requests'] == 'descendants': 405 term_processes = descendants 406 term_pids = descendant_pids 407 else: 408 # Process term requests by pgm_names. 409 term_processes = [] 410 for pgm_name in term_options['term_requests']['pgm_names']: 411 term_processes.extend(select_processes_by_pgm_name(descendants, pgm_name)) 412 term_pids = [str(process.pid) for process in term_processes] 413 414 message += gp.sprint_timen("Processes to be terminated:") \ 415 + gp.sprint_var(term_pids) 416 for process in term_processes: 417 process.terminate() 418 message += gp.sprint_timen("Waiting on the following pids: " + ' '.join(descendant_pids)) 419 gm.append_file(terminate_descendants_temp_file_path, message) 420 psutil.wait_procs(descendants) 421 422 # Checking after the fact to see whether any descendant processes are still alive. If so, a process 423 # report showing this will be included in the output. 424 descendants, descendant_pids, process_report = get_descendant_info(current_process) 425 if descendants: 426 message = "\n" + gp.sprint_timen("Not all of the processes terminated:") \ 427 + process_report 428 gm.append_file(terminate_descendants_temp_file_path, message) 429 430 message = gp.sprint_dashes(width=120) 431 gm.append_file(terminate_descendants_temp_file_path, message) 432 gp.print_file(terminate_descendants_temp_file_path) 433 os.remove(terminate_descendants_temp_file_path) 434 435 436def gen_exit_function(): 437 r""" 438 Execute whenever the program ends normally or with the signals that we catch (i.e. TERM, INT). 439 """ 440 441 # ignore_err influences the way shell_cmd processes errors. Since we're doing exit processing, we don't 442 # want to stop the program due to a shell_cmd failure. 443 ignore_err = 1 444 445 if psutil_imported and term_options: 446 terminate_descendants() 447 448 # Call the main module's exit_function if it is defined. 449 exit_function = getattr(module, "exit_function", None) 450 if exit_function: 451 exit_function() 452 453 gp.qprint_pgm_footer() 454 455 456def gen_signal_handler(signal_number, 457 frame): 458 r""" 459 Handle signals. Without a function to catch a SIGTERM or SIGINT, the program would terminate immediately 460 with return code 143 and without calling the exit_function. 461 """ 462 463 # The convention is to set up exit_function with atexit.register() so there is no need to explicitly 464 # call exit_function from here. 465 466 gp.qprint_executing() 467 468 # Calling exit prevents control from returning to the code that was running when the signal was received. 469 exit(0) 470 471 472def gen_post_validation(exit_function=None, 473 signal_handler=None): 474 r""" 475 Do generic post-validation processing. By "post", we mean that this is to be called from a validation 476 function after the caller has done any validation desired. If the calling program passes exit_function 477 and signal_handler parms, this function will register them. In other words, it will make the 478 signal_handler functions get called for SIGINT and SIGTERM and will make the exit_function function run 479 prior to the termination of the program. 480 481 Description of arguments: 482 exit_function A function object pointing to the caller's exit function. This defaults 483 to this module's gen_exit_function. 484 signal_handler A function object pointing to the caller's signal_handler function. This 485 defaults to this module's gen_signal_handler. 486 """ 487 488 # Get defaults. 489 exit_function = exit_function or gen_exit_function 490 signal_handler = signal_handler or gen_signal_handler 491 492 atexit.register(exit_function) 493 signal.signal(signal.SIGINT, signal_handler) 494 signal.signal(signal.SIGTERM, signal_handler) 495 496 497def gen_setup(): 498 r""" 499 Do general setup for a program. 500 """ 501 502 # Set exit_on_error for gen_valid functions. 503 gv.set_exit_on_error(True) 504 505 # Get main module variable values. 506 parser = getattr(module, "parser") 507 stock_list = getattr(module, "stock_list") 508 validate_parms = getattr(module, "validate_parms", None) 509 510 gen_get_options(parser, stock_list) 511 512 if validate_parms: 513 validate_parms() 514 sync_args() 515 gen_post_validation() 516 517 gp.qprint_pgm_header() 518