###############################################################
#
# Scriptname:  PRTGCustomCitrixRDSHCapacity.ps1
#
# Autor:       Urs Heeb
# Date:        13.09.2022
#
# Version:     2022.09.01 / 15.09.22 / Urs Heeb
#                  Create script based on other custom PRTG scripts
#                  
#
# Description: Script does following:
#              Connects to a Citrix controller
#              Gets the information of all delivery groups
#              Calculates the average load index of each delivery group 
#
# Requirements:
#              PRTG variables are needed while configuration
#              %host %windowsdomain %windowsuser %windowspassword
#              PRTG service user needs read permission on Citrix controller
#
###############################################################

###
# Get parameter from PRTG
param (
[string]$server,
[string]$domain,
[string]$username,
[string]$password
)

# For troubleshooting
<#
$server="controller.domain.pit"
$username="username"
$password="password"
#>

# Prepare credentials
$credentials = New-Object System.Management.Automation.PSCredential (($domain+'\'+$username), (ConvertTo-SecureString $password -AsPlainText -Force))

###
# Prepare an array of all delivery groups from a Citrix controller by an Invoke-Command
# Citrix PS commands are only available on those servers
$Data = $null
$Data=@()

$Data += Invoke-Command -Computername $server -credential $credentials -ScriptBlock {

    # Load Citrix PS Snapin
    add-pssnapin citrix*

    # Get delivery groups with the needed information
    $DeliveryGroups=Get-BrokerDesktopGroup | Select-Object Name, InMaintenanceMode

    # Set the array for later work
    $DG=@()
    ForEach ($group in $DeliveryGroups){
       $DG += [PSCustomObject]@{Name=$group.Name;InMaintenance=$group.InMaintenanceMode} 
    }

    # Return data to PRTG probe
    return $DG
}


###
# Prepare an array of the load of all delivery groups from a Citrix controller by an Invoke-Command
# Citrix PS commands are only available on those servers
$Load = $null
$Load=@()

ForEach ($RDSHGrp in $Data) {

    $RDSHArg = $null
    $RDSHArg = $RDSHGrp.Name

    # Determine if a delivery group is in maintenance
    # DGs in maintenance goes into warning state, but doesn't affect the whole sensor
    # DGs in maintenance will not be stated based on her usage
    # DGs in maintenance will have fake values of 9% and 99 to avoid an alarm
    If ($RDSHGrp.InMaintenance) {
        $RDSHWorkers=99
        $RDSHTotalLoad=99
        $RDSHAvgLoad=99
        $RDSHAvgPercent=9
    }
    Else {
        $RDSHLoad = $null
        $RDSHLoad = @()


        $RDSHScriptBlock =
            {
                 # Load Citrix PS Snapin
                add-pssnapin citrix*

                # Get delivery groups with the needed information
                $RDSHWorker=Get-BrokerMachine | Where-Object DesktopGroupName -eq $args[0] | Select-Object DNSName, LoadIndex

                # Set the array for later work
                $RDSH=@()
                ForEach ($Worker in $RDSHWorker){
                   $RDSH += [PSCustomObject]@{Name=$Worker.DNSName;Load=$Worker.LoadIndex} 
                }

                # Return data to PRTG probe
                return $RDSH
             }

        # Get the load index of all workers in the delivery group
        $RDSHLoad += Invoke-Command -Computername $server -credential $credentials -ScriptBlock $RDSHScriptBlock -ArgumentList $RDSHArg

        # Count the workers
        $RDSHWorkers = $RDSHLoad.count
        $RDSHTotalLoad = 0
        # Calculate the sum of all load indexes
        $RDSHLoad.Load | ForEach{$RDSHTotalLoad += $_}
        # Calculate the average load index for the delivery group
        $RDSHAvgLoad = "{0:F0}" -f ($RDSHTotalLoad/$RDSHWorkers)
        # Calculate the average percent value of the load index for the delivery group
        $RDSHAvgPercent = "{0:F0}" -f ($RDSHAvgLoad/100)

    }
    $Load += [PSCustomObject]@{Name=$RDSHArg;Worker=$RDSHWorkers;TotalLoad=$RDSHTotalLoad;AvgLoad=$RDSHAvgLoad;AvgPercent=$RDSHAvgPercent}
}


###
# Prepare the data and the PRTG XML output

# Reset some variables and define standard values
$returnState=$null
$returnState=@()
$returnStateOK = 0
$returnStateWarning = 1
$returnStateCritical = 2
$AlertLevel=90
$AlertString = "ALERT - high load index"
$WarningLevel=80
$WarningString = "Warning - load index reaches a high level"

# Start preparing XML output
$retXml = "<prtg>`n"

# Get data for each delivery group from the array created on the Citrix Controller
ForEach ($Dataset in $Load){
    # Prepare variables based on the array data
    $DGName = $Dataset.Name
    $DGTotal = $Dataset.Worker
    $DGLoad = $Dataset.AvgLoad
    # Calculate the percentage of the RDSH delivery groups
    $DGLoadPercent = $Dataset.AvgPercent
    $DGMaintenance = $Dataset.Maintenance

    # Determine return state if delivery group is not in maintenance
    # The values depend on the size of the delivery group
    If ([Int64]$DGLoadPercent -gt $AlertLevel){
        $RetState = $returnStateCritical
        $returnState += [PSCustomObject]@{Name=$DGName;State=[Int64]$RetState} 
    }
    ElseIf ([Int64]$DGLoadPercent -le $WarningLevel) {
        $RetState = $returnStateOK
        $returnState += [PSCustomObject]@{Name=$DGName;State=[Int64]$RetState} 
    }
    Else {
        $RetState = $returnStateWarning
        $returnState += [PSCustomObject]@{Name=$DGName;State=[Int64]$RetState} 
    }


    #$retXml += "  <result>`n"
    #$retXml += "    <channel>$DGName State</channel>`n"
    #$retXml += "    <value>$RetState</value>`n"
    #$retXml += "  </result>`n"
    $retXml += "  <result>`n"
    $retXml += "    <channel>$DGName Load %</channel>`n"
    $retXml += "    <value>$DGLoadPercent</value>`n"
    $retXml += "    <unit>Percent</unit>`n"
    $retXml += "    <limitMode>1</limitMode>`n"
    $retXml += "    <limitMaxError>$AlertLevel</limitMaxError>`n"
    $retXml += "    <limitErrorMsg>$AlertString</limitErrorMsg>`n"
    $retXml += "    <limitMaxWarning>$WarningLevel</limitMaxWarning>`n"
    $retXml += "    <limitWarningMsg>$WarningString</limitWarningMsg>`n"
    $retXml += "  </result>`n"
    $retXml += "  <result>`n"
    $retXml += "    <channel>$DGName Load Index</channel>`n"
    $retXml += "    <value>$DGLoad</value>`n"
    $retXml += "    <unit>Count</unit>`n"
    $retXml += "  </result>`n"
}

# Determine return string depends on the several states
If ($returnState.State -contains 2) {
    $RetString = $AlertString    
}
ElseIf ($returnState.State -contains 1) {
    $RetString = $WarningString    
}
Else {
    $RetString = "OK"    
}

$retXml += "  <text>$RetString</text>`n"
$retXml += "</prtg>`n"

###
# Return info to PRTG
write-host $retXml