#!/dataserfs/libs-2024-01-11/bin/python3

import json, os, paramiko, sys, time
from paramiko.client import AutoAddPolicy
from runbookpy import agent, utils

# Return 
def final_return(state, message):
    if state:
        print(json.dumps({'message': message, 'status': 'Success'}))
        sys.exit(0)
    else:
        print(json.dumps({'message': message, 'status': 'Failed'}))
        sys.exit(1)

class Hyperv:
    def __init__(self, hostname, username):
        self.hostname = hostname
        self.username = username
        self.password = os.getenv("WINDOWS_PASSWORD")

        if self.password == '':
            final_return(False, "Password is incorrect.")

    # Run remote SSH command
    def get_vm_details(self):
        # Get current time
        utc_ms = int(time.time()*1000)

        all_vm_details = []

        try:
            client = paramiko.SSHClient()
            client.set_missing_host_key_policy(AutoAddPolicy)
            client.connect(self.hostname, username=self.username, password=self.password)

            cmd = "Get-VM | ConvertTo-Json"
            
            stdin, rstdout, stderr = client.exec_command(f"PowerShell.exe -NoProfile -Command \"{cmd}\"")

            stdout = rstdout.read()

            get_vm_data = json.loads(stdout)

            for vm in get_vm_data:
                sub_cmd = f"Get-VMNetworkAdapter -VMName {vm['VMName']} | ConvertTo-Json"
                stdin, rstdout, stderr = client.exec_command(f"PowerShell.exe -NoProfile -Command \"{sub_cmd}\"")

                stdout = rstdout.read()

                vm_net_deets = json.loads(stdout)

                ## Actually format the data into a list
                ## colnames:
                ## name,state,version,status,type,subtype,notes,cpuusage,memoryassigned,memorydemand,autostartaction,autostopaction,
                ## computername,ticks,mac,dynamicmacenabled,islegacy,ip1,ip2,ip3,connected,switchname
                vm_details = []

                # append the data to the local list
                vm_details.append(vm['VMName'])
                vm_details.append(vm['State'])
                vm_details.append(vm['Version'])
                vm_details.append(vm['Status'])
                vm_details.append(vm['VirtualMachineType'])
                vm_details.append(vm['VirtualMachineSubType'])
                if "\n" in vm['Notes']:
                    notes = vm['Notes'].split('\n')[1]
                else:
                    notes = ""
                vm_details.append(notes)
                vm_details.append(vm['CPUUsage'])
                vm_details.append(vm['MemoryAssigned'])
                vm_details.append(vm['MemoryDemand'])
                vm_details.append(vm['AutomaticStartAction'])
                vm_details.append(vm['AutomaticStopAction'])
                vm_details.append(vm['ComputerName'])
                vm_details.append(vm['Uptime']['Ticks'])
                vm_details.append(vm_net_deets['MacAddress'])
                vm_details.append(vm_net_deets['DynamicMacAddressEnabled'])
                vm_details.append(vm_net_deets['IsLegacy'])

                # Handle varying amounts of IPs (3 seems to be the max, so work with that)
                if len(vm_net_deets['IPAddresses']) == 0:
                    vm_details.append("")
                    vm_details.append("")
                    vm_details.append("")
                elif len(vm_net_deets['IPAddresses']) == 1:
                    vm_details.append(vm_net_deets['IPAddresses'][0])
                    vm_details.append("")
                    vm_details.append("")
                elif len(vm_net_deets['IPAddresses']) == 2:
                    vm_details.append(vm_net_deets['IPAddresses'][0])
                    vm_details.append(vm_net_deets['IPAddresses'][1])
                    vm_details.append("")
                elif len(vm_net_deets['IPAddresses']) == 3:
                    vm_details.append(vm_net_deets['IPAddresses'][0])
                    vm_details.append(vm_net_deets['IPAddresses'][1])
                    vm_details.append(vm_net_deets['IPAddresses'][2])
                else:
                    # Super-old legacy
                    vm_details.extend(["", "", ""])

                vm_details.append(vm_net_deets['Connected'])
                vm_details.append(vm_net_deets['SwitchName'])

                # Append the list to the main list
                all_vm_details.append((vm_details))


            # Close SSH session
            client.close()

            # Clean up garbage that paramiko fails to do...
            del client, stdin, rstdout, stderr

            duration_ms = int(time.time()*1000) - utc_ms

            tsv = utils.list_to_tsv(all_vm_details)

            colnames = "name,state,version,status,type,subtype,notes,cpuusage,memoryassigned,memorydemand,autostartaction,autostopaction,computername,ticks,mac,dynamicmacenabled,islegacy,ip1,ip2,ip3,connected,switchname"

            print(self.to_rb_tsv("GetHyperVInfo", utc_ms, tsv, duration_ms, colnames))

        # In the event of local process return errors, fail out and remove DB entries
        except paramiko.ssh_exception.AuthenticationException:
            final_return(False, "Failed to authenticate on host.")        

    def to_rb_tsv(self, name, utc_ms, tsv, duration_ms, colnames):
        return agent.make_tsv("cmd", name, utc_ms, srcname="powershell", \
                    content=tsv, duration_ms=duration_ms, \
                    parserok="on", method="hyperv-monitor", colnames=colnames)

if __name__ == "__main__":
    if len(sys.argv) < 4:
        final_return(False, "Missing CLI Arguments.")

    # Define system
    ip      = sys.argv[1]
    user    = sys.argv[2]
    command = sys.argv[3]

    if len(sys.argv) > 4:
        command = sys.argv[3:]
        command = " ".join(command)
    else:
        command = sys.argv[3]

    # Initialize powershell class
    hyperv = Hyperv(ip, user)

    # Check commands
    if command == "get VM":
        hyperv.get_vm_details()

    else:
        final_return(False, "Invalid command!")

