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 209term_options = None 210 211 212def set_term_options(**kwargs): 213 r""" 214 Set the global term_options. 215 216 If the global term_options is not None, gen_exit_function() will call terminate_descendants(). 217 218 Description of arguments(): 219 kwargs Supported keyword options follow: 220 term_requests Requests to terminate specified descendants of this program. The 221 following values for term_requests are supported: 222 children Terminate the direct children of this program. 223 descendants Terminate all descendants of this program. 224 <dictionary> A dictionary with support for the following keys: 225 pgm_names A list of program names which will be used to identify which descendant 226 processes should be terminated. 227 """ 228 229 global term_options 230 # Validation: 231 arg_names = list(kwargs.keys()) 232 gv.valid_list(arg_names, ['term_requests']) 233 if type(kwargs['term_requests']) is dict: 234 keys = list(kwargs['term_requests'].keys()) 235 gv.valid_list(keys, ['pgm_names']) 236 else: 237 gv.valid_value(kwargs['term_requests'], ['children', 'descendants']) 238 term_options = kwargs 239 240 241if psutil_imported: 242 def match_process_by_pgm_name(process, pgm_name): 243 r""" 244 Return True or False to indicate whether the process matches the program name. 245 246 Description of argument(s): 247 process A psutil process object such as the one returned by psutil.Process(). 248 pgm_name The name of a program to look for in the cmdline field of the process 249 object. 250 """ 251 252 # This function will examine elements 0 and 1 of the cmdline field of the process object. The 253 # following examples will illustrate the reasons for this: 254 255 # Example 1: Suppose a process was started like this: 256 257 # shell_cmd('python_pgm_template --quiet=0', fork=1) 258 259 # And then this function is called as follows: 260 261 # match_process_by_pgm_name(process, "python_pgm_template") 262 263 # The process object might contain the following for its cmdline field: 264 265 # cmdline: 266 # [0]: /usr/bin/python 267 # [1]: /my_path/python_pgm_template 268 # [2]: --quiet=0 269 270 # Because "python_pgm_template" is a python program, the python interpreter (e.g. "/usr/bin/python") 271 # will appear in entry 0 of cmdline and the python_pgm_template will appear in entry 1 (with a 272 # qualifying dir path). 273 274 # Example 2: Suppose a process was started like this: 275 276 # shell_cmd('sleep 5', fork=1) 277 278 # And then this function is called as follows: 279 280 # match_process_by_pgm_name(process, "sleep") 281 282 # The process object might contain the following for its cmdline field: 283 284 # cmdline: 285 # [0]: sleep 286 # [1]: 5 287 288 # Because "sleep" is a compiled executable, it will appear in entry 0. 289 290 optional_dir_path_regex = "(.*/)?" 291 cmdline = process.as_dict()['cmdline'] 292 return re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[0]) \ 293 or re.match(optional_dir_path_regex + pgm_name + '( |$)', cmdline[1]) 294 295 def select_processes_by_pgm_name(processes, pgm_name): 296 r""" 297 Select the processes that match pgm_name and return the result as a list of process objects. 298 299 Description of argument(s): 300 processes A list of psutil process objects such as the one returned by 301 psutil.Process(). 302 pgm_name The name of a program to look for in the cmdline field of each process 303 object. 304 """ 305 306 return [process for process in processes if match_process_by_pgm_name(process, pgm_name)] 307 308 def sprint_process_report(pids): 309 r""" 310 Create a process report for the given pids and return it as a string. 311 312 Description of argument(s): 313 pids A list of process IDs for processes to be included in the report. 314 """ 315 report = "\n" 316 cmd_buf = "echo ; ps wwo user,pgrp,pid,ppid,lstart,cmd --forest " + ' '.join(pids) 317 report += gp.sprint_issuing(cmd_buf) 318 rc, outbuf = gc.shell_cmd(cmd_buf, quiet=1) 319 report += outbuf + "\n" 320 321 return report 322 323 def get_descendant_info(process=psutil.Process()): 324 r""" 325 Get info about the descendants of the given process and return as a tuple of descendants, 326 descendant_pids and process_report. 327 328 descendants will be a list of process objects. descendant_pids will be a list of pids (in str form) 329 and process_report will be a report produced by a call to sprint_process_report(). 330 331 Description of argument(s): 332 process A psutil process object such as the one returned by psutil.Process(). 333 """ 334 descendants = process.children(recursive=True) 335 descendant_pids = [str(process.pid) for process in descendants] 336 if descendants: 337 process_report = sprint_process_report([str(process.pid)] + descendant_pids) 338 else: 339 process_report = "" 340 return descendants, descendant_pids, process_report 341 342 def terminate_descendants(): 343 r""" 344 Terminate descendants of the current process according to the requirements layed out in global 345 term_options variable. 346 347 Note: If term_options is not null, gen_exit_function() will automatically call this function. 348 349 When this function gets called, descendant processes may be running and may be printing to the same 350 stdout stream being used by this process. If this function writes directly to stdout, its output can 351 be interspersed with any output generated by descendant processes. This makes it very difficult to 352 interpret the output. In order solve this problem, the activity of this process will be logged to a 353 temporary file. After descendant processes have been terminated successfully, the temporary file 354 will be printed to stdout and then deleted. However, if this function should fail to complete (i.e. 355 get hung waiting for descendants to terminate gracefully), the temporary file will not be deleted and 356 can be used by the developer for debugging. If no descendant processes are found, this function will 357 return before creating the temporary file. 358 359 Note that a general principal being observed here is that each process is responsible for the 360 children it produces. 361 """ 362 363 message = "\n" + gp.sprint_dashes(width=120) \ 364 + gp.sprint_executing() + "\n" 365 366 current_process = psutil.Process() 367 368 descendants, descendant_pids, process_report = get_descendant_info(current_process) 369 if not descendants: 370 # If there are no descendants, then we have nothing to do. 371 return 372 373 terminate_descendants_temp_file_path = gm.create_temp_file_path() 374 gp.print_vars(terminate_descendants_temp_file_path) 375 376 message += gp.sprint_varx("pgm_name", gp.pgm_name) \ 377 + gp.sprint_vars(term_options) \ 378 + process_report 379 380 # Process the termination requests: 381 if term_options['term_requests'] == 'children': 382 term_processes = current_process.children(recursive=False) 383 term_pids = [str(process.pid) for process in term_processes] 384 elif term_options['term_requests'] == 'descendants': 385 term_processes = descendants 386 term_pids = descendant_pids 387 else: 388 # Process term requests by pgm_names. 389 term_processes = [] 390 for pgm_name in term_options['term_requests']['pgm_names']: 391 term_processes.extend(select_processes_by_pgm_name(descendants, pgm_name)) 392 term_pids = [str(process.pid) for process in term_processes] 393 394 message += gp.sprint_timen("Processes to be terminated:") \ 395 + gp.sprint_var(term_pids) 396 for process in term_processes: 397 process.terminate() 398 message += gp.sprint_timen("Waiting on the following pids: " + ' '.join(descendant_pids)) 399 gm.append_file(terminate_descendants_temp_file_path, message) 400 psutil.wait_procs(descendants) 401 402 # Checking after the fact to see whether any descendant processes are still alive. If so, a process 403 # report showing this will be included in the output. 404 descendants, descendant_pids, process_report = get_descendant_info(current_process) 405 if descendants: 406 message = "\n" + gp.sprint_timen("Not all of the processes terminated:") \ 407 + process_report 408 gm.append_file(terminate_descendants_temp_file_path, message) 409 410 message = gp.sprint_dashes(width=120) 411 gm.append_file(terminate_descendants_temp_file_path, message) 412 gp.print_file(terminate_descendants_temp_file_path) 413 os.remove(terminate_descendants_temp_file_path) 414 415 416def gen_exit_function(): 417 r""" 418 Execute whenever the program ends normally or with the signals that we catch (i.e. TERM, INT). 419 """ 420 421 # ignore_err influences the way shell_cmd processes errors. Since we're doing exit processing, we don't 422 # want to stop the program due to a shell_cmd failure. 423 ignore_err = 1 424 425 if psutil_imported and term_options: 426 terminate_descendants() 427 428 # Call the main module's exit_function if it is defined. 429 exit_function = getattr(module, "exit_function", None) 430 if exit_function: 431 exit_function() 432 433 gp.qprint_pgm_footer() 434 435 436def gen_signal_handler(signal_number, 437 frame): 438 r""" 439 Handle signals. Without a function to catch a SIGTERM or SIGINT, the program would terminate immediately 440 with return code 143 and without calling the exit_function. 441 """ 442 443 # The convention is to set up exit_function with atexit.register() so there is no need to explicitly 444 # call exit_function from here. 445 446 gp.qprint_executing() 447 448 # Calling exit prevents control from returning to the code that was running when the signal was received. 449 exit(0) 450 451 452def gen_post_validation(exit_function=None, 453 signal_handler=None): 454 r""" 455 Do generic post-validation processing. By "post", we mean that this is to be called from a validation 456 function after the caller has done any validation desired. If the calling program passes exit_function 457 and signal_handler parms, this function will register them. In other words, it will make the 458 signal_handler functions get called for SIGINT and SIGTERM and will make the exit_function function run 459 prior to the termination of the program. 460 461 Description of arguments: 462 exit_function A function object pointing to the caller's exit function. This defaults 463 to this module's gen_exit_function. 464 signal_handler A function object pointing to the caller's signal_handler function. This 465 defaults to this module's gen_signal_handler. 466 """ 467 468 # Get defaults. 469 exit_function = exit_function or gen_exit_function 470 signal_handler = signal_handler or gen_signal_handler 471 472 atexit.register(exit_function) 473 signal.signal(signal.SIGINT, signal_handler) 474 signal.signal(signal.SIGTERM, signal_handler) 475 476 477def gen_setup(): 478 r""" 479 Do general setup for a program. 480 """ 481 482 # Set exit_on_error for gen_valid functions. 483 gv.set_exit_on_error(True) 484 485 # Get main module variable values. 486 parser = getattr(module, "parser") 487 stock_list = getattr(module, "stock_list") 488 validate_parms = getattr(module, "validate_parms", None) 489 490 gen_get_options(parser, stock_list) 491 492 if validate_parms: 493 validate_parms() 494 gen_post_validation() 495 496 gp.qprint_pgm_header() 497