import os
from typing import Dict

import docker
import pwn

from vagd import templates, helper
from vagd.box import Box
from vagd.virts.pwngd import Pwngd
from vagd.virts.shgd import Shgd


class Dogd(Shgd):
    """
    | Docker virtualization for pwntools

    :param binary: binary to execute
    :param image: docker base image
    :param user: name of user on docker container
    :param forward: Dictionary of forwarded ports, needs to follow docker api format: 'hostport/(tcp|udp)' : guestport
    :param ex: if experimental features, e.g. alpine, gdbserver should be enabled
    :param fast: mounts libs locally for faster symbol extraction (experimental) NOT COMPATIBLE WITH ALPINE
    :param kwargs: parameters to pass through to super

    | SSH from cmd

    .. code-block:: bash

        ssh -o "StrictHostKeyChecking=no" -i ~/.vagd/keyfile -p $(cut .vagd/docker.lock -d":" -f 2) vagd@0.0.0.0

    | connect with docker exec

    .. code-block:: bash

       docker exec -it $(cut ./.vagd/docker.lock -d":" -f 1) /bin/bash

    | Kill from cmd:

    .. code-block:: bash

        docker kill $(cut ./.vagd/docker.lock -d":" -f 1)

    | Docker containers are automatically removed after they stop
    | Docker images need to be manually removed from docker
    | Dockerfiles are stored in home directory to allow caching ~/.vagd/docker/<image>/Dockerfile

    .. code-block:: bash

        docker images # list images
        docker rmi <id> # remove correct image
    """

    _image: str
    _user: str
    _port: int
    _client: docker.client
    _id: str
    _dockerdir: str
    _dockerfile: str
    _isalpine: bool
    _gdbsrvport: int
    _ex: bool
    _forward: Dict[str, int]

    DOCKERHOME = Pwngd.HOME_DIR + "docker/"
    DEFAULT_USER = 'vagd'
    DEFAULT_PORT = 2222
    DEFAULT_IMAGE = Box.DOCKER_FOCAL

    DEFAULT_PACKAGES = Pwngd.DEFAULT_PACKAGES + ["openssh-server"]
    LOCKFILE = Pwngd.LOCAL_DIR + 'docker.lock'

    def _create_dockerfile(self):

        pwn.log.info(f'create new Dockerfile at f{self._dockerfile}')
        if not os.path.exists(Pwngd.KEYFILE):
            helper.generate_keypair()

        if not os.path.exists(self._dockerdir + "keyfile.pub"):
            os.link(Pwngd.PUBKEYFILE, self._dockerdir + "keyfile.pub")
        template = templates.DOCKER_ALPINE_TEMPLATE if self._isalpine else templates.DOCKER_TEMPLATE

        with open(self._dockerfile, 'w') as dockerfile:
            dockerfile.write(
                template.format(image=self._image,
                                packages=' '.join(Dogd.DEFAULT_PACKAGES),
                                user=self._user,
                                keyfile=os.path.basename(self._dockerdir + "keyfile.pub")))

    def _create_docker_instance(self):
        pwn.log.info('starting docker instance')
        self._port = helper.first_free_port(Dogd.DEFAULT_PORT)
        self._forward.update({'22/tcp': self._port})
        if self._isalpine:
            self._gdbsrvport = helper.first_free_port(Pwngd.STATIC_GDBSRV_PORT)
            self._forward.update({f'{self._gdbsrvport}/tcp': self._gdbsrvport})

        dir = os.path.dirname(os.path.realpath(__file__))
        with open(dir[:dir.rfind('/')] + '/res/seccomp.json', 'r') as seccomp_file:
            seccomp_rules = seccomp_file.read().strip()

        container = self._client.containers.run(self._bimage, ports=self._forward, detach=True, remove=True, security_opt=[f'seccomp:{seccomp_rules}'])
        self._id = container.id
        pwn.log.info(f'started docker instance {container.short_id}')
        with open(Dogd.LOCKFILE, 'w') as lockfile:
            lockfile.write(':'.join((container.id, str(self._port), str(self._gdbsrvport))))

    def _build_image(self):
        pwn.log.info('building docker image')
        return self._client.images.build(path=os.path.dirname(self._dockerfile), tag=f'vagd/{self._image}')[0]

    def _vm_create(self):

        if not os.path.exists(Pwngd.LOCAL_DIR):
            os.makedirs(Pwngd.LOCAL_DIR)

        if not os.path.exists(self._dockerfile):
            self._create_dockerfile()

        self._bimage = self._build_image()

        self._create_docker_instance()

    def _vm_setup(self) -> None:
        self._client = docker.from_env()
        if not os.path.exists(Dogd.LOCKFILE):
            pwn.log.info(f'No Lockfile {Dogd.LOCKFILE} found, creating new Docker Instance')
            self._vm_create()
        else:
            with open(Dogd.LOCKFILE, 'r') as lockfile:
                data = lockfile.readline().split(':')
                self._id = data[0]
                self._port = int(data[1])
                if self._isalpine:
                    self._gdbsrvport = int(data[2])
            if not self._client.containers.list(filters={'id':self._id}):
                pwn.log.info(f'Lockfile {Dogd.LOCKFILE} found, container not running, creating new one')
                self._vm_create()
            else:
                pwn.log.info(
                    f'Lockfile {Dogd.LOCKFILE} found, Docker Instance f{self._client.containers.get(self._id).short_id}')

    def __init__(self,
                 binary: str,
                 image: str = DEFAULT_IMAGE,
                 user: str = DEFAULT_USER,
                 forward: Dict[str, int] = None,
                 ex: bool = False,
                 fast: bool = False,
                 **kwargs):
        """

        :param binary: binary to execute
        :param image: docker base image
        :param user: name of user on docker container
        :param forward: Dictionary of forwarded ports, needs to follow docker api format: 'hostport/(tcp|udp)' : guestport
        :param ex: if experimental features, e.g. alpine, gdbserver should be enabled
        :param fast: mounts libs locally for faster symbol extraction (experimental) NOT COMPATIBLE WITH ALPINE
        :param kwargs: parameters to pass through to super
        """

        self._image = image
        self._isalpine = 'alpine' in image
        self._gdbsrvport = None
        self._dockerdir = Dogd.DOCKERHOME + f'{self._image}/'
        if not (os.path.exists(Dogd.DOCKERHOME) and os.path.exists(self._dockerdir)):
            os.makedirs(self._dockerdir)
        self._dockerfile = self._dockerdir + 'Dockerfile'
        self._user = user
        self._forward = forward
        self._ex = ex
        if self._isalpine and not self._ex:
            pwn.log.error("Docker alpine images requires experimental features")
        if self._forward is None:
           self._forward = dict()

        self._vm_setup()

        super().__init__(binary=binary,
                         user=self._user,
                         port=self._port,
                         ex=ex,
                         fast=fast,
                         gdbsrvport=self._gdbsrvport,
                         **kwargs)
