
import asyncio
import codecs
import confluent.util as util
import confluent.messages as msg
import confluent.tasks as tasks
import struct
import ssl

import aiohmi.util.webclient as webclient
import aiohmi.exceptions as pygexc
import confluent.interface.console as conapi
import io


class RetainedIO(io.BytesIO):
    # Need to retain buffer after close
    def __init__(self):
        self.resultbuffer = None
    def close(self):
        self.resultbuffer = self.getbuffer()
        super().close()

def fixuuid(baduuid):
    # VMWare changes the endian stuff in BIOS
    uuidprefix = (baduuid[:8], baduuid[9:13], baduuid[14:18])
    a = codecs.encode(struct.pack('<IHH', *[int(x, 16) for x in uuidprefix]),
        'hex')
    a = util.stringify(a)
    uuid = (a[:8], a[8:12], a[12:16], baduuid[19:23], baduuid[24:])
    return '-'.join(uuid).lower()

class VmConsole(conapi.Console):
    def __init__(self, host, port, tls, configmanager=None):
        self.tls = tls
        self.host = host
        self.port = port
        self.reader = None
        self.writer = None
        self.nodeconfig = configmanager
        self.connected = False
        self.recvr = None

    async def connect(self, callback):
        try:
            if self.tls:
                if not self.nodeconfig:
                    raise Exception('config manager instance required for TLS operation')
                kv = util.TLSCertVerifier(
                    self.nodeconfig, self.host, 'pubkeys.tls').verify_cert
                sslctx = ssl.create_default_context()
                sslctx.check_hostname = False
                sslctx.verify_mode = ssl.CERT_NONE
                self.reader, self.writer = await asyncio.open_connection(
                    self.host, self.port, ssl=sslctx)
                # The above disables default validation, replaced by
                # known-hosts style cert validation below
                sslobj = self.writer.get_extra_info('ssl_object')
                bincert = sslobj.getpeercert(binary_form=True)
                if not kv(bincert):
                    self.writer.close()
                    await self.writer.wait_closed()
                    raise pygexc.UnrecognizedCertificate('Unknown certificate', bincert)
            else:
                self.reader, self.writer = await asyncio.open_connection(
                    self.host, self.port)
        except Exception:
            await callback(conapi.ConsoleEvent.Disconnect)
            return
        self.connected = True
        self.datacallback = callback
        self.recvr = tasks.spawn_task(self.recvdata())

    async def write(self, data):
        self.writer.write(data)
        await self.writer.drain()

    async def close(self):
        self.connected = False
        if self.recvr:
            self.recvr.cancel()
            self.recvr = None
        if self.writer:
            self.writer.close()
            await self.writer.wait_closed()

    async def recvdata(self):
        try:
            while self.connected:
                try:
                    pendingdata = await self.reader.read(1024)
                except Exception:
                    pendingdata = b''
                if pendingdata == b'':
                    self.connected = False
                    await self.datacallback(conapi.ConsoleEvent.Disconnect)
                    return
                reply = b''
                while pendingdata and pendingdata[0] == 255:
                    cmd = pendingdata[1]
                    if cmd == 255:
                        pendingdata = pendingdata[1:]
                        break
                    subcmd = pendingdata[2]
                    if cmd == 253:  # DO
                        # binary, suppress go ohaed
                        if subcmd in (0, 3):
                            reply += b'\xff\xfb' + bytes([subcmd])  # will
                        else:
                            reply += b'\xff\xfc' + bytes([subcmd]) # won't do anything else
                        pendingdata = pendingdata[3:]
                    elif cmd == 251:  # will
                        # binary, suppress go ahead, echo
                        if subcmd in (0, 1, 3):
                            reply += b'\xff\xfd' + bytes([subcmd])  # do the implemented things
                        else:
                            reply += B'\xff\xfe' + bytes([subcmd])  # don't do others'
                        pendingdata = pendingdata[3:]
                    else:
                        raise Exception(repr(pendingdata[:3]))
                if reply:
                    await self.write(reply)
                if pendingdata:
                    await self.datacallback(pendingdata)
        except asyncio.CancelledError:
            pass



class VmwApiClient:
    def __init__(self, vcsa, user, password, configmanager):
        self.cachedurls = {}
        self.user = user
        self.password = password
        if configmanager:
            cv = util.TLSCertVerifier(
                configmanager, vcsa, 'pubkeys.tls_hardwaremanager'
            ).verify_cert
        else:
            def cv(x):
                return True

        try:
            self.user = self.user.decode()
            self.password = self.password.decode()
        except Exception:
            pass
        self.wc = webclient.WebConnection(vcsa, port=443, verifycallback=cv)
        self.logged = False
        self.vmlist = {}
        self.vmbyid = {}

    async def login(self):
        self.wc.set_basic_credentials(self.user, self.password)
        rsp = await self.wc.grab_json_response_with_status('/api/session', method='POST')
        body = rsp[0]
        del self.wc.stdheaders['Authorization']
        self.wc.set_header('vmware-api-session-id', body)
        self.logged = True

    async def get_screenshot(self, vm, outfile):
        if not self.logged:
            await self.login()
        vm = await self.index_vm(vm)
        url = f'/screen?id={vm}'
        self.wc.set_basic_credentials(self.user, self.password)
        try:
            data, status, headers = await self.wc.grab_response_with_status(url)
            if status == 200:
                outfile.write(data)
        finally:
            del self.wc.stdheaders['Authorization']

    async def list_vms(self):
        if not self.logged:
            await self.login()
        rsp = await self.wc.grab_json_response('/api/vcenter/vm')
        self.vmlist = {}
        for vm in rsp:
            name = vm['name']
            vmid = vm['vm']
            self.vmlist[name] = vmid
            self.vmbyid[vmid] = name
        return self.vmlist

    async def index_vm(self, vm):
        if not self.logged:
            await self.login()
        if vm in self.vmlist:
            return self.vmlist[vm]
        if vm in self.vmbyid:
            return vm
        await self.list_vms()
        if vm not in self.vmlist:
            if vm in self.vmbyid:
                return vm
            raise Exception("VM not found")
        return self.vmlist[vm]

    async def get_vm(self, vm):
        vm = await self.index_vm(vm)
        rsp = await self.wc.grab_json_response(f'/api/vcenter/vm/{vm}')
        return rsp

    async def get_vm_inventory(self, vm):
        rawinv = await self.get_vm(vm)
        hwver = rawinv['hardware']['version']
        uuid = fixuuid(rawinv['identity']['bios_uuid'])
        serial = rawinv['identity']['instance_uuid']
        sysinfo = {'name': 'System',
                   'present': True,
                   'information': {
                       'Product name': 'VMWare Virtual machine',
                       'UUID': uuid,
                       'Manufacturer': 'VMWare',
                       'Model': hwver,
                       'Serial Number': serial
                  }}
        inventory = [sysinfo]
        for nic in rawinv['nics']:
            nicinfo = rawinv['nics'][nic]
            label = nicinfo['label']
            mac = nicinfo['mac_address']
            inventory.append({
                'present': True,
                'name': label,
                'information': {
                'Type': 'Ethernet',
                'Model': nicinfo['type'],
                'MAC Address 1': mac}
                })
        yield msg.KeyValueData({'inventory': inventory}, vm)




    async def get_vm_host(self, vm):
        # unfortunately, the REST api doesn't manifest this as a simple attribute,
        vm = await self.index_vm(vm)
        rsp = await self.wc.grab_json_response('/api/vcenter/host')
        for hostinfo in rsp:
            host = hostinfo['host']
            rsp = await self.wc.grab_json_response(f'/api/vcenter/vm?hosts={host}')
            for guest in rsp:
                if guest['vm'] == vm:
                    return hostinfo

    async def get_vm_serial(self, vm):
        vm = await self.index_vm(vm)
        rsp = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/serial')
        if rsp[1] == 200 and len(rsp[0]) > 0:
            portid = rsp[0][0]['port']
            rsp = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/serial/{portid}')
            if rsp[1] == 200:
                if rsp[0]['backing']['type'] != 'NETWORK_SERVER':
                    return
                netloc = rsp[0]['backing']['network_location']
                portnum = netloc.split(':')[-1]
                tlsenabled = False
                if netloc.startswith('telnets'):
                    tlsenabled = True
                hostinfo = await self.get_vm_host(vm)
                hostname = hostinfo['name']
                rsp[0]
                return {
                    'server': hostname,
                    'port': portnum,
                    'tls': tlsenabled,
                    }

    async def get_vm_bootdev(self, vm):
        vm = await self.index_vm(vm)
        rsp = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/boot')
        if rsp[0]['enter_setup_mode']:
            return 'setup'
        rsp = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/boot/device')
        try:
            if rsp[0][0]['type'] == 'ETHERNET':
                return 'network'
        except IndexError:
            pass
        return 'default'

    async def get_vm_power(self, vm):
        vm = await self.index_vm(vm)
        rsp = await self.wc.grab_json_response(f'/api/vcenter/vm/{vm}/power')
        if rsp['state'] == 'POWERED_ON':
            return 'on'
        if rsp['state'] == 'POWERED_OFF':
            return 'off'
        if rsp['state'] == 'SUSPENDED':
            return 'suspended'
        raise Exception("Unknown response {}".format(repr(rsp)))

    async def set_vm_power(self, vm, state):
        current = None
        targstate = state
        vm = await self.index_vm(vm)
        if state == 'boot':
            current = await self.get_vm_power(vm)
            if current == 'on':
                state = 'reset'
                targstate = state
            else:
                targstate = 'on'
                state = 'start'
        elif state == 'on':
            state = 'start'
        elif state == 'off':
            state = 'stop'
        await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/power?action={state}', method='POST')
        return targstate, current


    async def set_vm_bootdev(self, vm, bootdev):
        vm = await self.index_vm(vm)
        self.wc.set_header('Content-Type', 'application/json')
        try:
            bootdevs = []
            entersetup = False
            if bootdev == 'setup':
                entersetup = True
            elif bootdev == 'default':
                # In theory, we should be able to send an empty device list.
                # However, vmware api counter to documentation seems to just ignore
                # such a request. So instead we just go "disk first"
                # and rely upon fast fail/retry to take us to a normal place
                currdisks, rcode = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/disk')
                currdisks = [x['disk'] for x in currdisks]
                bootdevs.append({'type': 'DISK', 'disks': currdisks})
            elif bootdev in ('net', 'network'):
                currnics, rcode = await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/ethernet')
                for nic in currnics:
                    bootdevs.append({'type': 'ETHERNET', 'nic': nic['nic']})
            payload = {'devices': bootdevs}
            await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/boot/device',
                                                        payload,
                                                        method='PUT')
            await self.wc.grab_json_response_with_status(f'/api/vcenter/vm/{vm}/hardware/boot',
                                                        {'enter_setup_mode': entersetup},
                                                        method='PATCH')
        finally:
            del self.wc.stdheaders['Content-Type']


def prep_vcsa_clients(nodes, configmanager):
    cfginfo = configmanager.get_node_attributes(nodes, ['hardwaremanagement.manager', 'secret.hardwaremanagementuser', 'secret.hardwaremanagementpassword'], decrypt=True)
    clientsbyvcsa = {}
    clientsbynode = {}
    for node in nodes:
        cfg = cfginfo[node]
        if 'hardwaremanagement.manager' not in cfg or 'value' not in cfg['hardwaremanagement.manager']:
            raise Exception("Missing hardwaremanagement.manager for node {}".format(node))
        currvcsa = cfg['hardwaremanagement.manager']['value']
        if currvcsa not in clientsbyvcsa:
             user = cfg.get('secret.hardwaremanagementuser', {}).get('value', None)
             passwd = cfg.get('secret.hardwaremanagementpassword', {}).get('value', None)
             clientsbyvcsa[currvcsa] = VmwApiClient(currvcsa, user, passwd, configmanager)
        clientsbynode[node] = clientsbyvcsa[currvcsa]
    return clientsbynode

async def retrieve(nodes, element, configmanager, inputdata):
    clientsbynode = prep_vcsa_clients(nodes, configmanager)
    for node in nodes:
        currclient = clientsbynode[node]
        if element == ['power', 'state']:
            yield msg.PowerState(node, await currclient.get_vm_power(node))
        elif element == ['boot', 'nextdevice']:
            yield msg.BootDevice(node, await currclient.get_vm_bootdev(node))
        elif element[:2] == ['inventory', 'hardware'] and len(element) == 4:
            async for rsp in currclient.get_vm_inventory(node):
                yield rsp
        elif element == ['console', 'ikvm_methods']:
            dsc = {'ikvm_methods': ['screenshot']}
            yield msg.KeyValueData(dsc, node)
        elif element == ['console', 'ikvm_screenshot']:
            # good background for the webui, and kitty
            imgdata = RetainedIO()
            imgformat = await currclient.get_screenshot(node, imgdata)
            imgdata = imgdata.getvalue()
            if imgdata:
                yield msg.ScreenShot(imgdata, node, imgformat=imgformat)
        elif element == ['health', 'hardware']:
            yield msg.HealthSummary('unknown', node)
            yield msg.SensorReadings([], node)







async def update(nodes, element, configmanager, inputdata):
    clientsbynode = prep_vcsa_clients(nodes, configmanager)
    for node in nodes:
        currclient = clientsbynode[node]
        if element == ['power', 'state']:
            newstate, oldstate = await currclient.set_vm_power(node, inputdata.powerstate(node))
            yield  msg.PowerState(node, newstate, oldstate)
        elif element == ['boot', 'nextdevice']:
            await currclient.set_vm_bootdev(node, inputdata.bootdevice(node))
            yield msg.BootDevice(node, await currclient.get_vm_bootdev(node))

# assume this is only console for now
async def create(nodes, element, configmanager, inputdata):
    clientsbynode = prep_vcsa_clients(nodes, configmanager)
    for node in nodes:
        serialdata = await clientsbynode[node].get_vm_serial(node)
        yield VmConsole(serialdata['server'], serialdata['port'], serialdata['tls'], configmanager)
        return



async def _selftest():
    import sys
    import os
    myuser = os.environ['VMWUSER']
    mypass = os.environ['VMWPASS']
    vc = VmwApiClient(sys.argv[1], myuser, mypass, None)
    vm = sys.argv[2]
    if sys.argv[3] == 'setboot':
        await vc.set_vm_bootdev(vm, sys.argv[4])
        await vc.get_vm_bootdev(vm)
    elif sys.argv[3] == 'power':
        await vc.set_vm_power(vm, sys.argv[4])
    elif sys.argv[3] == 'getinfo':
        await vc.get_vm(vm)
        print("Bootdev: " + await vc.get_vm_bootdev(vm))
        print("Power: " + await vc.get_vm_power(vm))
        print("Serial: " + repr(await vc.get_vm_serial(vm)))


if __name__ == '__main__':
    asyncio.run(_selftest())
