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