import os

from ...enums import CrashInputType
from .chess_exploit import CHESSExploit


class CHESSExploitControl(CHESSExploit):
    """
    An exploit that crashes with a controlled instruction pointer and register
    """

    def __init__(self, *args, registers=None, **kwargs):
        if registers is None:
            raise TypeError("Need to specify registers")
        super().__init__(*args, **kwargs)
        self.registers = registers

    def _write_script(self, **kwargs):
        if self.crash.input_type in (CrashInputType.STDIN, CrashInputType.POV_FILE):
            actions = self._script_get_actions_stdin()
        elif self.crash.input_type == CrashInputType.TCP:
            actions = self._script_get_actions_tcp()
        else:
            raise NotImplementedError("Unsupported crash input type %s." % self.crash.input_type)

        body = os.linesep.join(actions)

        preamble = """
import sys
import time
import nclib

if len(sys.argv) < 3:
    print("%s: <host> <port>" % sys.argv[0])
    sys.exit(1)

r = nclib.Netcat((sys.argv[1], int(sys.argv[2])), udp=False)
"""

        tail = "\nr.recvall()\n"
        for reg in self.registers:
            tail += 'print("REGISTER_%s=%X")\n' % (reg.upper(), self.registers[reg])

        return preamble + body + tail
