from weevely.core import messages
from weevely.core.loggers import log
from weevely.core.module import Module
from weevely.core.vectors import ModuleExec
from weevely.core.vectors import PhpCode
from weevely.core.vectors import ShellCmd


class Grep(Module):
    """Print lines matching a pattern in multiple files."""

    aliases = ["grep"]

    def init(self):
        self.register_info({"author": ["Emilio Pinna"], "license": "GPLv3"})

        # The grep action is done using multiple request.
        # First search for writable file, and then execute the grep
        # code for every found file. This allows to reuse code in
        # find_perm module and reduce the chances of timeout.

        self.register_vectors(
            [
                ShellCmd(
                    payload="grep ${ '' if case else '-i' } ${ '-v' if invert else '' } -e '${regex}' '${rfile}'",
                    name="grep_sh",
                    arguments=[
                        "-stderr_redirection",
                        " 2>/dev/null",
                    ],
                ),
                PhpCode(
                    payload=r"""% if invert:
$m=file_get_contents("${rfile}");$a=preg_replace("/${'' if regex.startswith('^') else '.*' }${regex.replace('/','\/')}${'' if regex.endswith('$') else '.*' }".PHP_EOL."?/m${ '' if case else 'i'}","",$m);if($a)print($a);
% else:
$m=Array();preg_match_all("/${'' if regex.startswith('^') else '.*' }${regex.replace('/','\/')}${'' if regex.endswith('$') else '.*' }/m${ '' if case else 'i'}",file_get_contents('${rfile}'),$m);if($m) print(implode(PHP_EOL,$m[0]));
% endif""",
                    name="grep_php",
                ),
            ]
        )

        self.register_arguments(
            [
                {"name": "rpath", "help": "Path. If is a folder grep all the contained files."},
                {"name": "regex", "help": "Regular expression to match file name"},
                {"name": "-case", "help": "Search case sensitive expression", "action": "store_true", "default": False},
                {"name": "-name-regex", "help": "Regular expression to match file name to grep"},
                {"name": "-no-recursion", "action": "store_true", "default": False},
                {"name": "-output", "help": "Redirect output to remote file"},
                {
                    "name": "-v",
                    "dest": "invert",
                    "action": "store_true",
                    "default": False,
                    "help": "Invert matching to select non-matching lines",
                },
                {"name": "-local", "action": "store_true", "default": False, "help": "Save redirected output locally"},
                {"name": "-vector", "choices": self.vectors.get_names(), "default": "grep_php"},
            ]
        )

    def run(self, **kwargs):
        files = []

        if ModuleExec("file_check", [self.args["rpath"], "dir"]).run():
            # If remote path is a folder, harvest all the readable
            # files wih given name-regex

            # Prepare the arguments for file_find
            file_find_args = ["-readable", self.args["rpath"], "-ftype", "f"]
            if self.args.get("name_regex"):
                file_find_args += ["-name-regex", self.args.get("name_regex")]
            if self.args.get("no_recursion"):
                file_find_args += ["-no-recursion"]

            files = ModuleExec("file_find", file_find_args).run()

        elif (
            ModuleExec("file_check", [self.args["rpath"], "file"]).run()
            and ModuleExec("file_check", [self.args["rpath"], "readable"]).run()
        ):
            # If the remote path is a readable file, just store the path
            files = [self.args["rpath"]]

        # Validate files presence
        if not isinstance(files, list) or not files:
            log.warning(messages.module_file_grep.failed_retrieve_info)
            return None, False

        # Store the found data in data dictionary in the
        # form `{ filename : [ line1, line2, ... ] }`
        # and store them as string whether requested
        results = {}
        output_str = ""
        output_path = self.args.get("output")

        for rfile in files:
            result_str = self.vectors.get_result(
                self.args["vector"],
                {"regex": self.args["regex"], "rfile": rfile, "case": self.args["case"], "invert": self.args["invert"]},
            )

            # Both grep_sh and grep_php -v trail a \n, this should work fine
            result_str = result_str[:-1] if result_str and result_str.endswith("\n") else result_str

            result_list = result_str.split("\n") if isinstance(result_str, str) and result_str else []

            # This means the command returned something something
            if result_list:
                if output_path:
                    # If output is redirected, just append to output_str
                    output_str += result_str

                # Else, print it out
                elif len(files) > 1:
                    # Print filepath:line if there are multiple files

                    for line in result_list:
                        log.info(f"{rfile}:{line}")
                else:
                    # Else, just print the lines
                    log.info("\n".join(result_list))

                results[rfile] = result_list

        # Save output to file whether specified
        saved = False
        if output_path:
            if not self.args.get("local"):
                saved = ModuleExec("file_upload", ["-content", result_str, output_path]).run()
            else:
                try:
                    with open(output_path, "wb") as outputfile:
                        outputfile.write(result_str.encode("utf-8"))
                except Exception as e:
                    log.warning(messages.generic.error_loading_file_s_s % (output_path, str(e)))
                else:
                    saved = True

        return results, saved

    def print_result(self, result):
        pass