import sys
import socket
import time
import array
import paramiko
import xml.etree.ElementTree as ET


class gmp:
    gmv_hostname = '127.0.0.1'
    gvm_port = 22
    gvm_username = 'omp'
    gvm_password = ''
    gvm_timeout = 5

    gvm_gmp_sock = ''
    gvm_gmp_stdin = ''
    gvm_gmp_stdout = ''

    def __init__(self):
        pass

    def connect(self):
        pass

    def connect(self, hostname, port, username, password, timeout):
        self.gvm_hostname = hostname
        self.gvm_port = port
        self.gvm_username = username
        self.gvm_password = password
        self.gvm_timeout = timeout

        self.gvm_gmp_sock = paramiko.SSHClient()
        self.gvm_gmp_sock.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            self.gvm_gmp_sock.connect(hostname=hostname, username=username, password=password, timeout=timeout, port=port)
        except (paramiko.BadHostKeyException, paramiko.AuthenticationException, paramiko.SSHException, socket.error) as e:
            print('SSH Connection failed: ' + str(e))
            sys.exit(1)

        channel = self.gvm_gmp_sock.invoke_shell()

        self.gvm_gmp_stdin = channel.makefile('wb')
        self.gvm_gmp_stdout = channel.makefile('r')

        self.gvm_gmp_stdin.write('/usr/bin/socat UNIX-CONNECT:/usr/local/var/run/openvasmd.sock -\n')
        time.sleep(0.5)
        self.gvm_gmp_stdout.read(len(self.gvm_gmp_stdout.channel.in_buffer))

        print(self.gvm_hostname
              + ' ' + str(self.gvm_port)
              + ' ' + self.gvm_username
              + ' ' + self.gvm_password
              + ' ' + str(self.gvm_timeout))

    def close(self):
        self.gvm_gmp_sock.close()

    def authenticate(self, username='admin', password='eebd0396-6aae-4c03-94ae-5d1bd53da656'):
        cmd = '<authenticate><credentials><username>' + username + '</username><password>' + password + '</password></credentials></authenticate>'
        self.sendCommand(cmd)

    def getTasks(self):
        self.sendCommand('<get_tasks/>')

    def sendCommand(self, cmd):
        #print('About to send this here: ' + cmd)
        self.gvm_gmp_stdin.write(str(cmd) + '\n')
        time.sleep(0.2)
        buf_len = len(self.gvm_gmp_stdout.channel.in_buffer)
        data_len = len(cmd) + 2
        raw_xml_data = self.gvm_gmp_stdout.read(buf_len).decode()[data_len:]
        print(raw_xml_data)

        decoded_data = array
        root = ET.fromstring(raw_xml_data)

        #return ET.tostringlist(root)
        for child in root.findall('.'):
            # if there are children, get their text content as well.
            if len(child):
                for subchild in child:
                    print(subchild.text)
            # else just get the current child text.
            else:
                print(child.text)





import sys
import argparse
import gmp
import code

def main(argv):
   # hostname = 'localhost'
   # username = ''
    quick_xml_command = ''

    gvm = gmp.gmp()
    gvm.connect(hostname='172.30.10.188', port=22, username='omp', password='', timeout=5)
    #test.connect(hostname='172.30.10.188', port=22, username='omp', password='', timeout=5)
    #print(argv)

    if  args.xml and len(args.xml) != 0:
        # Start execute single xml command
        startSingleMode(gvm, args.xml)
    elif args.interactive:
        # Start the interactive Shell
        startInteractiveMode(gvm)

    input_data = ''
    while True:
        input_data = input('gvm-shell: ')
        if input_data == 'quit':
            break
        gvm.sendCommand(input_data)

    gvm.close()

def startSingleMode(gvm, xml_command):
    gvm.sendCommand(xml_command)
    gvm.close()
    sys.exit(0)

def startInteractiveMode(gvm):
    code.interact(banner='GVM Interactive Console', local=dict(globals(), **locals()))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='gvm-shell')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-X', '--xml', help='Send a xml command to OpenVAS Manager Socket')
    group.add_argument('-i', '--interactive', action='store_true', default=False, help='Start an interactive shell')
    parser.add_argument('-I', '--input', help='Read a pre written gmp file. It will be executed after successful connection')

    args = parser.parse_args()

    main(args)