__author__ = '@tomereyz'

import argparse
import os

DOCKERFILE_TEMPLATE = """

{architecture_dependant}
RUN apt-get update

WORKDIR /

RUN apt-get install -y openssh-server
RUN apt-get install -y sudo
RUN apt-get install -y python
RUN apt-get install -y gdb
RUN apt-get install -y git
RUN apt-get install -y vim
RUN apt-get install -y gcc
RUN apt-get install -y strace

{challenge_template}

RUN mkdir /var/run/sshd
RUN echo 'root:default' | chpasswd

RUN sed -i 's/[#]*PermitRootLogin [a-zA-Z\-]*/PermitRootLogin yes/' /etc/ssh/sshd_config
RUN sed -i 's/[#]*UsePrivilegeSeparation [a-zA-Z\-]*/UsePrivilegeSeparation no/' /etc/ssh/sshd_config
RUN sed -i 's/[#]*Banner [a-zA-Z\-\/\.]*/Banner \/etc\/banner/' /etc/ssh/sshd_config

RUN sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
RUN echo "export VISIBLE=now" >> /etc/profile
RUN sysctl -w kernel.dmesg_restrict=1
RUN chmod 1733 /tmp /var/tmp /dev/shm

RUN git clone https://github.com/longld/peda.git /usr/bin/peda
RUN echo "source /usr/bin/peda/peda.py" >> ~/.gdbinit
RUN chmod 0755 /usr/bin/peda/*.py
RUN chmod 0755 /usr/bin/peda/lib/*.py

RUN sysctl -w kernel.randomize_va_space={aslr}


EXPOSE 22
CMD ["/usr/sbin/sshd", "-D"]
"""

CHALLENGE_TEMPLATE = """

COPY {architecture_challenges}/{challenge_name}/* /home/{challenge_name}/

RUN adduser {challenge_name} ; adduser {challenge_name}_root ; usermod -G {challenge_name} {challenge_name} ; usermod -G {challenge_name}_root {challenge_name}_root ; /home/{challenge_name}/compile.sh /home/{challenge_name}

RUN chown {challenge_name}_root /home/{challenge_name}/flag ; chown {challenge_name}_root:{challenge_name} /home/{challenge_name}/{challenge_name} ; chown {challenge_name}:{challenge_name} /home/{challenge_name}/{challenge_name}.c /home/{challenge_name}/compile.sh ; chown root:root /home/{challenge_name}_root

RUN chmod 0400 /home/{challenge_name}/flag ; chmod 4550 /home/{challenge_name}/{challenge_name} ; chmod 0440 /home/{challenge_name}/{challenge_name}.c /home/{challenge_name}/compile.sh

RUN echo '{challenge_name}:{challenge_name}' | chpasswd ; echo '{challenge_name}_root:default' | chpasswd ; echo "source /usr/bin/peda/peda.py" > /home/{challenge_name}/.gdbinit

"""


class DockerFile(object):
    def __init__(self, sources, docker_name, no_cache, port, build_run, arch, aslr):
        """
        Generate dockerfile
        :param sources:
        :param docker_name:
        :param no_cache:
        :param port:
        :param build_run:
        """
        self.sources = sources
        self.dockerfile = None
        self.docker_name = docker_name
        self.no_cache = '--no-cache' if no_cache else ''
        self.port = port
        self.build_run = build_run
        self.arch = arch
        self.aslr = True if aslr == 'yes' else False

    def __enter__(self):
        self.generate()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return

    def _format_template(self):
        challenge_template = ''
        if self.arch == 'i386':
            architecture_challenges = "exercises-i386-aslr" if self.aslr else 'exercises-i386'
            architecture_dependant = "FROM i386/ubuntu\nCOPY {arch_c}/banner /etc/banner\nCOPY {arch_c}/motd /etc/motd".format(
                arch_c=architecture_challenges)
        else:
            architecture_challenges = "exercises-arm32v7-aslr" if self.aslr else 'exercises-arm32v7'
            architecture_dependant = "FROM armhf/ubuntu\nCOPY qemu-arm-static /usr/bin/\nCOPY {arch_c}/motd /etc/motd\nCOPY {arch_c}/banner /etc/banner".format(
                arch_c=architecture_challenges)

        for f in self.sources:
            challenge_template += CHALLENGE_TEMPLATE.format(challenge_name=f.split('.c')[0],
                                                            architecture_challenges=architecture_challenges)

        dockerfile = DOCKERFILE_TEMPLATE.format(architecture_dependant=architecture_dependant,
                                                challenge_template=challenge_template,
                                                aslr='2' if self.aslr else '0')

        self.dockerfile = dockerfile

    def _docker_build_run(self):
        interactive = raw_input(
            'execute `sudo docker stop $(sudo docker ps -a -q) ; sudo docker rm $(sudo docker ps -a -q)` y/n?')
        if interactive == 'y':
            os.system('sudo docker stop $(sudo docker ps -a -q) ; sudo docker rm $(sudo docker ps -a -q)')
        os.system('sudo docker build {no_cache} -t {docker_name} .'.format(no_cache=self.no_cache,
                                                                           docker_name=self.docker_name))
        os.system(
            'sudo docker run --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -d -p {port}:22 --name {docker_name} {docker_name}'.format(
                port=self.port, docker_name=self.docker_name))

    def generate(self):
        self._format_template()
        with open('Dockerfile', 'w') as f:
            f.write(self.dockerfile)

        if self.build_run:
            self._docker_build_run()


def main(port, arch, sources, docker_name, aslr, no_cache=True, build_run=False):
    with DockerFile(sources=sources, docker_name=docker_name, port=port, no_cache=no_cache,
                    build_run=build_run, arch=arch, aslr=aslr) as handle:
        print handle.sources


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Scan directory and prepares challenges & Flags')
    parser.add_argument('-s', '--sources', help='challenges sources names', required=True, dest='sources')
    parser.add_argument('-d', '--docker-name', help='docker name', required=True, dest='docker_name')
    parser.add_argument('-ca', '--no-cache', help='use cache when building docker', required=False, dest='no_cache',
                        action='store_false')
    parser.add_argument('-r', '--build-run', help='build and run docker', required=False, dest='build_run',
                        action='store_true')
    parser.add_argument('-p', '--port', help='running docker port', required=True, dest='port')

    parser.add_argument('-a', '--architecture', help='i386/arm32v7', required=True, dest='arch')
    parser.add_argument('-as', '--aslr', help='yes/no', required=True, dest='aslr')

    main(**vars(parser.parse_args()))
