import telnetlib

from v0ltlib.utils.v0lt_utils import fail, green, yellow, bytes_to_str


# WIP - SOME FEATURES NEED SOME REWORK


class Telnet:
    tn = None

    def __init__(self, hostname, port):
        try:
            self.tn = telnetlib.Telnet(hostname, port)
        except Exception as err:
            fail("Could not connect Telnet to {0}:{1} (error: {2})".format(hostname, port, err))
        print("Connected to port {0}".format(green(port)))

    def write(self, command):
        self.tn.write(bytes(command, "UTF-8"))

    def writeln(self, command):
        self.write(command + "\n")

    def shellnet(self, shellcode):
        shellcode = shellcode.replace("\\x", "")
        self.tn.write(bytearray.fromhex(shellcode))

    def shellnetln(self, shellcode):
        self.shellnet(shellcode + "\n")

    def read(self, nb_of_recv=1):
        data = "\n"
        for x in range(0, nb_of_recv):
            data += bytes_to_str(self.tn.read_some())
        return data

    def read_until(self, substring):
        return bytes.decode(self.tn.read_until(bytes(substring, "UTF-8")), "UTF-8")

    def dialogue(self, command, nb_of_recv=1):
        self.writeln(command)
        return "{0}: {1}".format(yellow("Answer"), self.read(nb_of_recv))

@@ -156,60 +156,256 @@
     
-tn = None%0D%0A%0D%0A    def __init__(self, hostname, port):
+'''%0D%0A    Quick Telnet implementation%0D%0A    '''%0D%0A%0D%0A    def __init__(self, hostname, port):%0D%0A        '''%0D%0A        Initialize a netcat client%0D%0A%0D%0A        :param hostname: Hostname to connect to%0D%0A        :param port:     Port to connect to%0D%0A        '''%0D%0A
 %0D%0A  
@@ -689,32 +689,84 @@
 elf, command):%0D%0A
+        '''%0D%0A        Send a message%0D%0A        '''%0D%0A%0D%0A
         self.tn.
@@ -824,32 +824,101 @@
 elf, command):%0D%0A
+        '''%0D%0A        Send a message followed by a %5Cn%0D%0A        '''%0D%0A%0D%0A
         self.wri
@@ -967,32 +967,108 @@
 f, shellcode):%0D%0A
+        '''%0D%0A        Send a shellcode from the ShellCrafter%0D%0A        '''%0D%0A%0D%0A
         shellcod
@@ -1168,114 +1168,195 @@
 def 
-shellnetln(self, shellcode):%0D%0A        self.shellnet(shellcode + %22%5Cn%22)%0D%0A%0D%0A    def read(self, nb_of_recv=1):
+read(self, nb_of_recv=1):%0D%0A        '''%0D%0A        Read from the server%0D%0A        :param nb_of_recv: Number of message to try to receive%0D%0A        :returns:    The received data%0D%0A        '''%0D%0A
 %0D%0A  
@@ -1491,25 +1491,24 @@
  data%0D%0A%0D%0A   
- 
 def read_unt
@@ -1541,133 +1541,582 @@
     
-return bytes.decode(self.tn.read_until(bytes(substring, %22UTF-8%22)), %22UTF-8%22)%0D%0A%0D%0A    def dialogue(self, command, nb_of_recv=1):
+'''%0D%0A        Read until 'substring' is found%0D%0A%0D%0A        :param substring: The string to look for%0D%0A        :returns:         The received data%0D%0A        '''%0D%0A%0D%0A        return bytes.decode(self.tn.read_until(bytes(substring, %22UTF-8%22)), %22UTF-8%22)%0D%0A%0D%0A    def dialogue(self, command, nb_of_recv=1):%0D%0A        '''%0D%0A        Exchange with the server%0D%0A        Sends 'commands' and waits for 'nb_of_recv' messages back%0D%0A%0D%0A        :param command:    message to send%0D%0A        :param nb_of_recv: Number of messages to try to read%0D%0A        :returns:          The received data%0D%0A        '''%0D%0A
 %0D%0A  

Add docstring to telnet