import os

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 PhpFile
from weevely.utils.strings import str2hex


class Curl(Module):
    """Perform a curl-like HTTP request."""

    aliases = ["curl"]

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

        self.register_vectors(
            [
                PhpFile(
                    payload_path=os.path.join(self.folder, "php_context.tpl"),
                    name="file_get_contents",
                    arguments=["-raw-response"],
                ),
                PhpFile(
                    payload_path=os.path.join(self.folder, "php_context.tpl"),
                    name="fopen_stream_get_contents",
                    arguments=["-raw-response"],
                ),
                PhpFile(
                    payload_path=os.path.join(self.folder, "php_context.tpl"),
                    name="fopen_fread",
                    arguments=["-raw-response"],
                ),
                PhpFile(
                    payload_path=os.path.join(self.folder, "php_curl.tpl"), name="php_curl", arguments=["-raw-response"]
                ),
                PhpFile(
                    payload_path=os.path.join(self.folder, "php_httprequest1.tpl"),
                    name="php_httprequest1",
                    arguments=["-raw-response"],
                ),
                # TODO: fix this, it fails the "POST request with binary string" test
                # due to some bash limitation with null bytes.
                # ShellCmd(
                #   payload = """curl -s -i ${ '-A "$(env echo -ne \"%s\")"' % user_agent if user_agent else "" } ${ '--connect-timeout %i' % connect_timeout } ${ '-X %s' % request if (not data and request) else '' } ${ " ".join([ '-H "$(env echo -ne \"%s\")"' % h for h in header ]) } ${ '-b "$(env echo -ne \"%s\")"' % cookie if cookie else '' } ${ '--data-binary $(env echo -ne "%s")' % ' '.join(data) if data else '' } ${ '$(env echo -ne "%s")' % url }""",
                #   name = 'sh_curl'
                # )
            ]
        )

        self.register_arguments(
            [
                {"name": "url"},
                {"name": "--header", "dest": "header", "action": "append", "default": []},
                {"name": "-H", "dest": "header", "action": "append", "default": []},
                {"name": "--cookie", "dest": "cookie"},
                {"name": "-b", "dest": "cookie"},
                {"name": "--data", "dest": "data", "action": "append", "default": []},
                {"name": "-d", "dest": "data", "action": "append", "default": []},
                {"name": "--user-agent", "dest": "user_agent"},
                {"name": "-A", "dest": "user_agent"},
                {"name": "--connect-timeout", "type": int, "default": 5, "help": "Default: 2"},
                {
                    "name": "--request",
                    "dest": "request",
                    "choices": ("GET", "HEAD", "POST", "PUT", "OPTIONS"),
                    "default": "GET",
                },
                {
                    "name": "-X",
                    "dest": "request",
                    "choices": ("GET", "HEAD", "POST", "PUT", "OPTIONS"),
                    "default": "GET",
                },
                {"name": "--output", "dest": "output"},
                {"name": "-o", "dest": "output"},
                {
                    "name": "-i",
                    "dest": "include_headers",
                    "help": "Include response headers",
                    "action": "store_true",
                    "default": False,
                },
                {
                    "name": "-local",
                    "action": "store_true",
                    "default": False,
                    "help": "Save file locally with -o|--output",
                },
                {"name": "-vector", "choices": self.vectors.get_names(), "default": "file_get_contents"},
            ]
        )

    def _encode(self):
        self.args["url"] = str2hex(self.args["url"])

        if self.args["data"]:
            self.args["data"] = [str2hex(x) for x in self.args["data"]]

        if self.args["user_agent"]:
            self.args["user_agent"] = str2hex(self.args["user_agent"])

        if self.args["cookie"]:
            self.args["cookie"] = str2hex(self.args["cookie"])

        if self.args["header"]:
            self.args["header"] = [str2hex(x) for x in self.args["header"]]

    def run(self, **kwargs):
        headers = []
        saved = None

        self._encode()

        vector_name, result = self.vectors.find_first_result(
            names=[self.args.get("vector")], format_args=self.args, condition=lambda r: r if r and r.strip() else None
        )

        # Print error and exit with no response or no headers
        if not (vector_name and result):
            log.warning(messages.module_net_curl.unexpected_response)
            return None, headers, saved
        if b"\r\n" * 2 not in result:
            # If something is returned but there is \r\n*2, we consider
            # everything as header. It happen with responses 204 No contents
            # that end with \r\n\r (wtf).
            headers = result
            result = b""
        else:
            headers, result = result.split(b"\r\n" * 2, 1)

        headers = [h.rstrip() for h in headers.split(b"\r\n")] if b"\r\n" in headers else headers

        output_path = self.args.get("output")
        if output_path:
            # If response must be saved, it's anyway safer to save it
            # within additional requests
            if not self.args.get("local"):
                saved = ModuleExec("file_upload", ["-content", result, output_path]).run()
            else:
                try:
                    with open(output_path, "wb") as resultfile:
                        resultfile.write(result)
                except Exception as e:
                    log.warning(messages.generic.error_loading_file_s_s % (output_path, str(e)))
                    saved = False
                else:
                    saved = True

        return result, headers, saved

    def print_result(self, result):
        resultstring = result[0].decode("utf-8", "replace")
        headers = [r.decode("utf-8", "replace") for r in result[1]]
        saved = result[2]

        # If is saved, we do not want output
        if self.args.get("output"):
            log.info(saved)
            return

        # Print headers if requested
        if self.args.get("include_headers"):
            log.info("\r\n".join(headers) + "\r\n")

        if resultstring:
            log.info(resultstring)