"""
FortiMonitor Countermeasure script helper - base class to allow easy setup of standalone scripts
to run as a countermeasure.

Copyright 2023 Fortinet, Inc. All Rights Reserved.
[email protected]


To use, create a subclass of CountermeasureScriptHelper and define the following properties:

 - name - A human-readable name for the countermeasure
 - textkey - A unique textkey describing the countermeasure
 - command: The command line to execute
 - description: Optional longer description of what the plugin does
 - capture_output: True or False value of whether to report the full output of the script

For example:

class TmpUsageCountermeasure(CountermeasureScriptHelper):
    name = "/tmp disk usage"
    textkey = "disk.tmp_usage"
    description = "Get the total usage of hte /tmp partition"
    command = "df -u /tmp"
    capture_output = True

"""

from CountermeasurePlugin import CountermeasurePlugin


class CountermeasureScriptHelper(CountermeasurePlugin):
    wall_announce_delay = None
    max_frequency = None
    max_runtime = None
    sudo_requirements = []
    author = "[email protected]"

    # The command to execute as part of the countermeasure - needs to be overridden in inheriting classes
    command = None

    # Whether to capture the output of the script and report as the result of the countermeasure
    capture_output = True

    def validate(self):
        problems = []
        if self.name == "Base Countermeasure":
            problems.append("Missing name definition")
        if self.textkey == "base":
            problems.append("Missing textkey definition")
        if self.command is None:
            problems.append("Missing command definition")
        if self.capture_output not in (True, False):
            problems.append("Invalid value for capture_output")

        return problems and ", ".join(problems) or None

    def run(self):
        if self.command is None:
            self.log.error(
                "No command specified for %s Countermeasure" % self.__class__.__name__
            )
            return

        return_code, output = self.execute(self.command)
        if self.capture_output:
            self.save_text_output(output)
        else:
            self.save_text_output(
                "Completed execution of %s Countermeasure" % self.__class__.__name__
            )
        self.save_return_code(return_code)