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