import re

from weevely.core import messages
from weevely.core.loggers import log
from weevely.core.module import Module
from weevely.core.module import Status
from weevely.core.vectors import PythonCode
from weevely.core.vectors import ShellCmd


class Su(Module):
    """Execute commands with su."""

    aliases = ["ifconfig"]

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

        self.register_vectors(
            [
                ShellCmd(
                    """expect -c 'spawn su -c "${command}" "${user}"; expect -re "assword"; send "${ passwd }\r\n"; expect eof;'""",
                    name="sh_expect",
                    postprocess=lambda x: re.findall("Password: (?:\r\n)?([\\s\\S]+)", x)[0]
                    if "Password: " in x
                    else "",
                ),
                PythonCode(
                    """
                    import pexpect as p,sys
                    c = p.spawn("su ${user} -c ${command}")
                    c.expect(".*assword:");c.sendline("${ passwd }")
                    i = c.expect([p.EOF,p.TIMEOUT])
                    if i!=p.TIMEOUT:
                        sys.stdout.write(c.before[3:].decode("utf-8","replace"))
                    """,
                    name="pyexpect",
                ),
            ]
        )

        self.register_arguments(
            [
                {"name": "passwd", "help": "User's password"},
                {"name": "command", "help": "Shell command", "nargs": "+"},
                {"name": "-user", "help": "User to run the command with", "default": "root"},
                {"name": "-stderr_redirection", "default": " 2>&1"},
                {
                    "name": "-vector-sh",
                    "choices": (
                        "system",
                        "passthru",
                        "shell_exec",
                        "exec",
                        "popen",
                        "proc_open",
                        "perl_system",
                        "pcntl",
                    ),
                },
                {"name": "-vector", "choices": self.vectors.get_names()},
            ]
        )

    def setup(self):
        """Probe all vectors to find a working su command.

        The method run_until is not used due to the check of shell_sh
        enabling for every tested vector.

        Args:
            self.args: The dictionary of arguments

        Returns:
            Status value, must be Status.RUN, Status.FAIL, or Status.IDLE.

        """

        args_check = {"user": self.args["user"], "passwd": self.args["passwd"], "command": "whoami"}

        (vector_name, result) = self.vectors.find_first_result(
            names=[self.args.get("vector", "")],
            format_args=args_check,
            condition=lambda result: (
                # Stop if shell_sh is in FAIL state
                self.session["shell_sh"]["status"] == Status.FAIL
                or
                # Or if the result is correct
                (self.session["shell_sh"]["status"] == Status.RUN and result and result.rstrip() == self.args["user"])
            ),
        )

        if self.session["shell_sh"]["status"] == Status.RUN and result and result.rstrip() == self.args["user"]:
            self.session["shell_su"]["stored_args"]["vector"] = vector_name
            return Status.RUN
        log.warn(messages.module_shell_su.error_su_executing)
        return Status.IDLE

    def run(self, **kwargs):
        # Join the command list and

        # Escape the single quotes. This does not protect from \' but
        # avoid to break the query for an unscaped quote.

        self.args["command"] = " ".join(self.args["command"]).replace("'", "\\'")

        format_args = {"user": self.args["user"], "passwd": self.args["passwd"], "command": self.args["command"]}

        if self.args.get("vector_sh"):
            format_args["vector"] = self.args["vector_sh"]

        if self.args.get("stderr_redirection"):
            format_args["stderr_redirection"] = self.args["stderr_redirection"]

        return self.vectors.get_result(name=self.args["vector"], format_args=format_args)