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

import json, os, paramiko, socket, 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 Powershell:
    def __init__(self, hostname, username):
        self.hostname = hostname
        self.username = username
        self.password = os.getenv("WINDOWS_PASSWORD")

        if self.password == '':
            final_return(False, "Password in correct.")

    # Run remote SSH command
    def run_ssh_cmd(self, cmd):
        try:
            client = paramiko.SSHClient()
            client.set_missing_host_key_policy(AutoAddPolicy)
            client.connect(self.hostname, username=self.username, password=self.password)
            
            stdin, rstdout, rstderr = client.exec_command(f"PowerShell.exe -NoProfile -Command \"{cmd}\"")

            stdout = rstdout.read()
            stderr = rstderr.read()
            
            client.close()

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

            # If there's anything in stderr, treat it as error output
            if stderr:
                return stderr.decode('utf-8').strip().splitlines()


            if not stdout:
                return None

            try:
                return json.loads(stdout)
            except json.JSONDecodeError:
                return stdout.decode('utf-8').strip().splitlines()

        # 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.")        

    # Get hotfix information
    def get_hotfix(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-HotFix | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames="hostname,description,hotfixid,installedby,installedon"
        hotfix_data = []

        # Parse and gather the pertinent data
        for line in data:
            # Append as list for list_to_tsv later
            hotfix_data.append([self.hostname, line['Description'], line['HotFixID'], line['InstalledBy'], line['InstalledOn']['DateTime']])

        # Convert list to TSV
        tsv = utils.list_to_tsv(hotfix_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetHotFix", utc_ms, tsv, duration_ms, colnames))

    # Get Physical Disks information
    def get_physical_disks(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-PhysicalDisk | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames="deviceid,name,mediatype,operationstatus,healthstatus,usage,size"
        physical_disks_data = []

        # If only 1 disk exists, will not return a list of dicts but a single dict, we need to convert to list
        if isinstance(data, dict):
            data = [data]

        # Parse and gather the pertinent data
        for line in data:
            # Append as list for list_to_tsv later
            physical_disks_data.append([(line['DeviceId']), line['FriendlyName'], line['MediaType'], line['OperationalStatus'], line['HealthStatus'], line['Usage'], str(line['Size'])])

        # Convert list to TSV
        tsv = utils.list_to_tsv(physical_disks_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetPhysicalDisks", utc_ms, tsv, duration_ms, colnames))

    # Get service list
    def get_service(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-Service | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)
        
        # Setup for RB TSV
        colnames = "status,name,displayname"
        service_data = []

        # Parse and gather the pertinent data 
        for line in data:
            # Parse status from int to string based off status #
            status = line['Status']
            if status == 1:
                status = "Stopped"
            elif status == 2:
                status = "StartPending"
            elif status == 3:
                status = "StopPending"
            elif status == 4:
                status = "Running"
            elif status == 5:
                status = "ContinuePending"
            elif status == 6:
                status = "PausePending"
            elif status == 7:
                status = "Paused"
            else:
                status = "Unknown"

            # Append as list for list_to_tsv later
            service_data.append([status, line['Name'], line['DisplayName']])

        # Convert the list to TSV
        tsv = utils.list_to_tsv(service_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetService", utc_ms, tsv, duration_ms, colnames))

    # Get SMB Share 
    def get_smb_share(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-SMBShare | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames = "name,scope,path,description"
        smb_share_data = []
        
        # Parse and gather pertinent data
        for line in data:
            # Append as list for list_to_tsv later
            smb_share_data.append([line['Name'], line['ScopeName'], line['Path'], line['Description']])

        # Convert list to TSV
        tsv = utils.list_to_tsv(smb_share_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetSmbShare", utc_ms, tsv, duration_ms, colnames))

    # Get computer system info
    def get_class_win32_compsys(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-WmiObject -Class win32_computersystem | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames = "domain,manufacturer,model,name,primaryownername,totalphysicalmemory"
        compsys_data = []

        # Parse and gather pertinent data
        compsys_data.append([data['Domain'], data['Manufacturer'], data['Model'], data['Name'], data['PrimaryOwnerName'], str(data['TotalPhysicalMemory'])])

        # Convert list to TSV
        tsv = utils.list_to_tsv(compsys_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetWmiObjectClassWin32CompSys", utc_ms, tsv, duration_ms, colnames))

    # Get operating System info
    def get_win32_opsys(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-WmiObject win32_OperatingSystem | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames = "buildnumber,registereduser,serialnumber,version"
        win32_opsys_data = []

        # Parse and gather pertinent data
        win32_opsys_data.append([data['BuildNumber'], data['RegisteredUser'], data['SerialNumber'], data['Version']])

        # Convert list to TSV
        tsv = utils.list_to_tsv(win32_opsys_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetWin32Opsys", utc_ms, tsv, duration_ms, colnames))

    # Get serial number and SMBIOS Asset Tag
    def get_win32_sysenclosure(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "Get-WmiObject Win32_SystemEnclosure | Select-Object SerialNumber,SMBIOSAssetTag | ConvertTo-Json"

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Setup for RB TSV
        colnames = "serialnumber,smbiosassettag"
        win32_sysenclosure_data = []

        # Parse and gather pertinent data
        win32_sysenclosure_data.append([data['SerialNumber'], data['SMBIOSAssetTag']])

        # Convert list to TSV
        tsv = utils.list_to_tsv(win32_sysenclosure_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetWin32Sysenclosure", utc_ms, tsv, duration_ms, colnames))

    def get_pending_updates(self):
        # Get current time
        utc_ms = int(time.time() * 1000)
        cmd = ("$Session = New-Object -ComObject Microsoft.Update.Session; "
            "$Searcher = $Session.CreateUpdateSearcher(); "
            "$Results = $Searcher.Search('IsInstalled=0 and Type=''Software'''); "
            "$Updates = foreach ($Update in $Results.Updates) { "
            "$kb = if ($Update.KBArticleIDs) { $Update.KBArticleIDs -join ',' } else { 'N/A' }; "
            "[PSCustomObject]@{ KB = $kb; Title = $Update.Title; IsDownloaded = $Update.IsDownloaded; "
            "IsInstalled = $Update.IsInstalled; RebootRequired = $Update.RebootRequired } }; "
            "if ($Updates) { $Updates | ConvertTo-Json -Compress } else { '[]' }")

        # Run remote SSH command
        data = self.run_ssh_cmd(cmd)

        # Determine status
        # Check for constrained language mode error
        if isinstance(data, list) and any("Cannot create type. Only core types are supported" in line for line in data):
            success = False
            status_msg = "Cannot check pending updates: PowerShell Constrained Language Mode does not support Windows Update COM objects."

        # If data is empty or not a list, treat as no updates
        elif not data or (isinstance(data, str) and data.strip() in ('', '[]')):
            success = True
            status_msg = "No pending updates found."

        # If data is a list and has at least one dict, treat as found; if data is a dict (single update), also treat as found
        elif (isinstance(data, list) and any(isinstance(line, dict) for line in data)) or isinstance(data, dict):
            success = True
            status_msg = "Pending updates found."

        # Otherwise, treat as no updates
        else:
            success = True
            status_msg = "No pending updates found."

        # Prepare TSV data
        colnames = "machine,success,status"
        tsv_data = [[self.hostname, str(success), status_msg]]

        # Convert list to TSV
        tsv = utils.list_to_tsv(tsv_data)

        # Calculate duration
        duration_ms = int(time.time() * 1000) - utc_ms

        # Return TSV using to_rb_tsv
        print(self.to_rb_tsv("GetPendingUpdates", utc_ms, tsv, duration_ms, colnames))
        
    def get_bitlocker_status(self):
        # Get current time
        utc_ms = int(time.time()*1000)
        cmd = "if (Get-Command Get-BitLockerVolume -ErrorAction SilentlyContinue) {     Get-BitLockerVolume | ConvertTo-Json } else { Write-Output \"Get-BitLockerVolume is not available.\" }"

        # Run remote SSH command (returns raw JSON string)
        data = self.run_ssh_cmd(cmd)

        if isinstance(data, list):
            if any(isinstance(line, str) and "Get-BitLockerVolume is not available." in line for line in data):
                final_return(True, "BitLocker not supported.")

        # Setup for RB TSV
        colnames="volumetype,mountpoint,capacitygb,metadataversion,volumestatus,encryptionpercentage,keyprotector,autolockenabled,protectionstatus,encryptionmethod"
        bitlocker_data = []

        # Parse and gather the pertinent data
        for line in data:
            if not isinstance(line, dict):
                continue  # skip lines that are not dicts
            bitlocker_data.append([
                line.get('VolumeType', ''),
                line.get('MountPoint', ''),
                str(line.get('CapacityGB', '')),
                str(line.get('MetadataVersion', '')),
                line.get('VolumeStatus', ''),
                str(line.get('EncryptionPercentage', '')),
                str(line.get('KeyProtector', '')),
                str(line.get('AutoLockEnabled', '')),
                str(line.get('ProtectionStatus', '')),
                line.get('EncryptionMethod', '')
            ])

        # If no valid BitLocker data, return final message
        if not bitlocker_data:
            final_return(True, "No BitLocker volumes found.")

        # Convert list to TSV
        tsv = utils.list_to_tsv(bitlocker_data)

        # Calculate duration
        duration_ms = int(time.time()*1000) - utc_ms

        # Print runbook TSV
        print(self.to_rb_tsv("GetBitLockerStatus", utc_ms, tsv, duration_ms, colnames))

    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="powershell-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
    powershell = Powershell(ip, user)

    # Check commands
    if command == "get hotfix":
        powershell.get_hotfix()

    elif command == "get physical disks":
        powershell.get_physical_disks()

    elif command == "get service":
        powershell.get_service()

    elif command == "get smb share":
        powershell.get_smb_share()

    elif command == "get class win32 compsys":
        powershell.get_class_win32_compsys()
    
    elif command == "get win32 opsys":
        powershell.get_win32_opsys()

    elif command == "get win32 sysenclosure":
        powershell.get_win32_sysenclosure()

    elif command == "get pending updates":
        powershell.get_pending_updates()

    elif command == "get bitlocker status":
        powershell.get_bitlocker_status()

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

