import importlib
import sys
import time
from abc import ABC, abstractmethod
from pathlib import Path

from pwnic import db
from pwnic.utils import install_package


class Exploit(ABC):
    def __init__(self, name):
        self.name = name
        # self.log  # implement logger based on object instance (can i use the name?)
        return

    def register(self):
        exploits[self.name] = self

    @abstractmethod
    def attack(self, target_ip: str, target: dict | None = None) -> list[str]:
        pass

    def run(
        self,
        target_ip: str,
        target: dict | list | None = None,
        targets_id: int | None = None,
    ):
        t = time.perf_counter()

        flags: list[str] = self.attack(target_ip, target)

        elapsed_time = time.perf_counter() - t

        db.saveExploitRun(
            self.name, t, elapsed_time, target_ip, flags, targets_id
        )


class PythonExploit(Exploit):
    def __init__(
        self,
        name,
        path: str | Path | None = None,
        dir: str | Path | None = None,
    ):
        """
        path is the path to the file to import and to consider as entrypoint.
        dir is the path of the directory containing said file.
        If a path is specified and is valid, dir is ignored.
        If the specified path is a dir and no dir is specified than it is used as such.
        if both a path and a dir were specified but path is a dir, if they are equal then the dir is considered. If they are not, an exception is raised

        If the specified path is invalid or if no path is specified then `dir` is scanned looking for known entrypoints in order: `__main__.py`, `main.py`, `__init__.py`. In none can be found an error is raised.
        Default dir is exploits/{name}
        """

        super().__init__(name)
        if path is not None:
            if not Path(path).is_dir():
                self.path = Path(path)
                self.dir = Path(path).parent
            else:
                if dir is None:
                    self.dir = Path(path)
                elif path == dir:
                    self.dir = Path(dir)
                else:
                    raise Exception
                self.path = self.findPath(self.dir)

        else:
            if dir is not None:
                self.dir = Path(dir)
            else:
                self.dir = Path(f"exploits/python/{name}")
            self.path = self.findPath(self.dir)

        self.check_imports()

    def findPath(self, dir: str | Path) -> Path:
        if not Path(dir).is_dir():
            raise Exception(f"{Path(dir)} is not a directory")
        files = list(Path(dir).glob("**/[!_]*.py"))
        if len(files) == 1:
            return files[0]
        elif "__main__.py" in files:
            return Path(dir, "__main__.py")
        elif "main.py" in files:
            return Path(dir, "main.py")
        elif "__init__.py" in files:
            return Path(dir, "__init__.py")
        else:
            raise Exception(
                f"{self.name} is defined as a python exploit but has more than one .py file and none is '__main__.py', 'main.py' or '__init__.py'. \
            Please define an entrypoint."
            )

    def check_imports(self):
        retries = 1
        while retries >= 0:
            try:
                spec = importlib.util.spec_from_file_location(
                    self.name, self.path.resolve()
                )
                module = importlib.util.module_from_spec(spec)  # type: ignore
                sys.modules[self.name] = module
                spec.loader.exec_module(module)  # type: ignore
                self.exploit_function: function = getattr(
                    sys.modules[self.name], "main"
                )
                return
            except ModuleNotFoundError as e:
                modulename = (
                    e.msg.split()[-1].replace("'", "").replace('"', "")
                )
                install_package(modulename)
            except AttributeError:
                raise
            except Exception:
                retries -= 1

    def attack(self, target_ip: str, target: dict | None = None) -> list[str]:
        flags: list[str] = self.exploit_function(target_ip, target)
        return flags


def load_exploit(type: str, name: str, *args, **kwargs) -> Exploit:
    if type == "python":
        exploit = PythonExploit(name, *args, **kwargs)
    else:
        raise ValueError(
            f"{type} is not a valid exploit type.",
        )

    exploit.register()
    return exploit


exploits: dict[str, Exploit] = {}
