<# 
.Synopsis 
    Gets the user/pool assignement details of virtual machines that are managed by the RD Connection Broker Server on which the script is being executed.
    
.Description 
    The script displays the user/pool assignement details of virtual machines that are managed by the RD Connection Broker Server on which the script is being executed. It displays the list of all the virtual machines and their associated status - assigned/unassigned to users/pools.
    
    Note: This script requires privileges to query the domain or trusted domains (if required) and hence must be run with the appropriate user credentials.
          To get a summary of the count of virtual machines, based on assignment type, refer to the examples.

.Parameter User
    Name(s) of the user(s) or user group(s) for whom the virtual machine assignment details are to be fetched.

.Parameter VirtualMachine
    Name(s) of the virtual machine(s) to get the assignment status.

.Parameter Pool
    Name(s) of the pool(s) to enumerate virtual machine(s) assigned to it.

.Parameter UserAssigned
    Switch: Displays information only for virtual machines that are assigned to users. If querying with "User" parameter, fetches virtual machine(s) information only for the specific user(s) or user group(s).

.Parameter PoolAssigned
    Switch: Displays information only for virtual machines that are assigned to pool(s).

.Parameter UnAssigned
    Switch: Displays information only for virtual machines that are currently unassigned. If querying with "User" parameter, fetches virtual machine(s) information assigned to all users except the ones specified.

.Example 
    PS C:\> .\Get-VMAssignment.ps1
    
    Displays the list of all the virtual machines managed by the RD Connection Broker server on which the script is executing and their assignment status.
    
.Example 
    PS C:\> .\Get-VMAssignment.ps1 -UserAssigned
    
    Displays the list of all the virtual machines managed by the RD Connection Broker server on which the script is executing that are assigned to users. It also displays the details of the user to which the virtual machine has been assigned.
    
.Example 
    PS C:\> .\Get-VMAssignment.ps1 -PoolAssigned
    
    Displays the list of all the virtual machines managed by the RD Connection Broker server on which the script is executing that are assigned to pools. It also displays the name of the pool to which the virtual machine has been assigned.

.Example 
    PS C:\> .\Get-VMAssignment.ps1 -UnAssigned
    
    Displays the list of all the virtual machines managed by the RD Connection Broker server on which the script is executing that are not assigned to users or pools.

.Example 
    PS C:\> .\Get-VMAssignment.ps1 -User mydomain\user01, trusteddomain\group01
    
    Displays the virtual machine(s) assignment status for the user "mydomain\user01" and all users belonging to the group "trusteddomain\group01".

.Example 
  
.Example 
    PS C:\> .\Get-VMAssignment.ps1 -User trusteddomain\group01 -UnAssigned
    
    Enumerates users belonging to the group "trusteddomain\group01" who do not have virtual machines assigned to them.

.Example 
    PS C:\> .\Get-VMAssignment.ps1 -VirtualMachine VM01.mydomain.myorg.com, VM02.mydomain.myorg.com
    
    Displays the assignment related information for the virtual machines VM01.mydomain.myorg.com and VM02.mydomain.myorg.com.

.Example 
    PS C:\> .\Get-VMAssignment.ps1 -Pool Win7Pool
    
    Displays all the virtual machines assigned to the pool Win7Pool.

.Example 
    PS C:\> .\Get-VMAssignment.ps1 | Group-Object -Property AssignmentType -NoElement
    
    Summarizes the virtual machine assignemnt details, displays the count of virtual machines that have been assigned. 
    
.Outputs 
    PSObject(s).
    These object(s) returned have the following information related to the virtual machine:
    
    Property        Description 
    --------------------------------------------------
    VM              Name of the Virtual Machine.
    Host            RD Virtualization host server on which this virtual machine is running.
    AssignedTo      Name of the user or pool to which this virtual machine is assigned.
    AssignmentType  Has values "User" or "Pool", based on whether the virtual machine is assigned to a user or a pool.

    These object(s) contain following information related to User assignment:
    
    Property        Description 
    --------------------------------------------------
    User            Name of the user.
    VM              Name of the Virtual Machine assigned to this user.
    Host            RD Virtualization host server on which this virtual machine is hosted.
    
.Link 
    Set-PersonalVirtualDesktop
    Clear-PersonalVirtualDesktop   
#>

[CmdletBinding(DefaultParametersetName="SimpleFilter")]
param (

    [Parameter(Mandatory=$TRUE, ParameterSetName="Users")]
    [String[]]
    $User,

    [Parameter(Mandatory=$TRUE, ParameterSetName="virtual machines")]
    [String[]]
    $VirtualMachine,

    [Parameter(Mandatory=$TRUE, ParameterSetName="Pools")]
    [String[]]
    $Pool,

    [Parameter(Mandatory=$FALSE, ParameterSetName="Users")]
    [Parameter(Mandatory=$FALSE, ParameterSetName="SimpleFilter")]
    [Switch]
    $UserAssigned,
    
    [Parameter(Mandatory=$FALSE, ParameterSetName="SimpleFilter")]
    [Switch]
    $PoolAssigned,
    
    [Parameter(Mandatory=$FALSE, ParameterSetName="Users")]
    [Parameter(Mandatory=$FALSE, ParameterSetName="SimpleFilter")]
    [Switch]
    $UnAssigned
)

$propertyNames = @{
    [System.DirectoryServices.ActiveDirectory.DomainMode]::Windows2008Domain = "mstsproperty01"
    [System.DirectoryServices.ActiveDirectory.DomainMode]::Windows2008R2Domain = "mstsproperty01"
}

function Get-FQDN ([string]$distinguishedName)
{
	$nameTranslator = New-Object -ComObject NameTranslate
	$translatorType = $nameTranslator.GetType()
	[void]$translatorType.InvokeMember("init", ”InvokeMethod”, $NULL, $nameTranslator, @(3, $NULL))
	[void]$translatorType.InvokeMember("set", ”InvokeMethod”, $NULL, $nameTranslator, @(1, $distinguishedName))

	$mName = $translatorType.InvokeMember("get", ”InvokeMethod”, $NULL, $nameTranslator, 3).split("\")[1].trim("$")
	$mDomain = $translatorType.InvokeMember("get", ”InvokeMethod”, $NULL, $nameTranslator, 2).split("/")[0]
	$fqdn = "{0}.{1}" -f $mName, $mDomain

	return $fqdn
}

function Get-DomainNetBIOS ([string]$domain)
{
    $rootDSE = New-Object System.DirectoryServices.DirectoryEntry "LDAP://$domain/rootDSE"
    if (-not $rootDSE.Properties) {return}
    $domainCN = $rootDSE.Properties['defaultNamingContext'][0]

    $parts = New-Object System.DirectoryServices.DirectoryEntry "LDAP://$domain/CN=Partitions,CN=Configuration,$domainCN"
    $parts.Children | ?{$_.Properties["nCName"][0] -eq $domainCN} | %{$_.Properties["NetBIOSName"][0]}
}

function Get-VMAssignmentDictionary ([string[]]$domainNames)
{
    $compUserDict = @{}
    
    $currentDomain = [System.DirectoryServices.ActiveDirectory.Domain]::GetComputerDomain()
    $allTrustedDomains = $currentDomain.GetAllTrustRelationships() | ?{$_.TrustDirection -eq "Outbound" -OR $_.TrustDirection -eq "Bidirectional"}

    $domainsToQuery = @($currentDomain.Name)
    $allTrustedDomains | ?{$_} | %{$domainsToQuery += $_.TargetName}

    $domainNetbiosMap = @{}
    $domainsToQuery | %{$domainNetbiosMap[$_] = (Get-DomainNetBIOS $_)}
    
    if (($domainNames -ne $NULL) -AND ($domainNames.Count -ge 1))
    {
        $domainNames | ?{$domainNetbiosMap.Values -notcontains $_} | %{Write-Warning ("Unknown domain : $_")}
        $domainsToQuery =  $domainNetbiosMap.Keys | ?{$domainNames -contains $domainNetbiosMap[$_]}
        
        Write-Verbose "Querying domains $($domainsToQuery -join ", ")"
    }
    
    foreach ($domain in $domainsToQuery)
    {
        $domainContext = New-Object System.DirectoryServices.ActiveDirectory.DirectoryContext 'Domain', $domain
        
        $domObject = $NULL
        try { $domObject = [System.DirectoryServices.ActiveDirectory.Domain]::GetDomain($domainContext) }
        catch {Write-Error "Unable to connect to domain : $domain. $($_.Exception.InnerException.Message)"}
        
        if (-not $domObject) { continue }
        
        $desktopPropertyName = $propertyNames[$domObject.DomainMode]
        if ($desktopPropertyName -eq $NULL) { $desktopPropertyName = "mstsproperty01" }
    
        $dirEntry = New-Object System.DirectoryServices.DirectoryEntry "LDAP://$domain"
        $objsearch = New-Object System.DirectoryServices.DirectorySearcher $dirEntry
        $objsearch.Filter = "(& (ObjectCategory=user)($desktopPropertyName=*))"
        [void]$objsearch.PropertiesToLoad.Add("cn")
        [void]$objsearch.PropertiesToLoad.Add("samaccountname")
        [void]$objsearch.PropertiesToLoad.Add($desktopPropertyName)
        
        foreach ($entry in $objsearch.FindAll())
        {
            if($entry.Properties[$desktopPropertyName] -ne $NULL)
            {
                $distinguishedName = $entry.Properties[$desktopPropertyName] | ?{$_ -like "machine=*"}
                $NULL, $friendlyVMName = $distinguishedName -split "="
                $friendlyUserName = "{0}\{1}" -f ($domainNetbiosMap[$domain]), $entry.Properties["samaccountname"][0]
                
                if ($friendlyVMName)
                {
                    $compUserDict[$friendlyVMName] = $friendlyUserName
                }
            }
        }
        
        $objsearch.Dispose()
    }
    
    return $compUserDict
}


function Convert-ToDomainUserFormat([string]$adsiPath)
{
    if ($adsiPath -match "WinNT://(.*)/(.*)")
    {
        return ("{0}\{1}" -f $matches[1..2])
    }
}

function Get-User([string]$adsiEntry)
{
    if ($adsiEntry -like "*\*")
    {
        $adsiPath = "WinNT://{0}/{1}" -f ($adsiEntry -split "\\")
    }
    else
    {
        $adsiPath = "WinNT://{0}/{1}" -f ((Get-WmiObject Win32_ComputerSystem).Domain.split('.')[0]), $adsiEntry
    }
    $adsiObj = [ADSI]$adsiPath
    
    if(-not $adsiObj.Name)
    {
        Write-Error "Invalid user/group : $adsiEntry"
        return
    }
    
    if ($adsiObj.Class -eq "User")
    {
        return (Convert-ToDomainUserFormat $adsiObj.Path)
    }
    
    $groupMembers = $adsiObj.Members() | %{[System.DirectoryServices.DirectoryEntry]$_}
    foreach ($member in $groupMembers)
    {
        switch($member.Class)
        {
            "Group" {Get-User (Convert-ToDomainUserFormat $member.Path)}
            "User"  {Write-Output (Convert-ToDomainUserFormat $member.Path)}
        }
    }
}



function Get-PoolAssignment()
{
    $vmList = @(Get-WmiObject -Namespace root\cimv2 -Query "Select * from Win32_SessionBrokerTarget where pluginname='VmResource' and farmname != ''")
    
    $poolDict = @{}
    Get-WmiObject -Namespace root\cimv2\terminalservices -Class Win32_TSRemoteDesktop | %{$poolDict[$_.Alias] = $_.Name}
    
    foreach ($vm in $vmList)
    {
        $vmEntry = New-Object PSObject

        $vmEntry | Add-Member -MemberType NoteProperty -Name VM -Value $vm.TargetName
        $vmEntry | Add-Member -MemberType NoteProperty -Name Host -Value $vm.Environment
        $vmEntry | Add-Member -MemberType NoteProperty -Name AssignedTo -Value $poolDict[$vm.FarmName]
        $vmEntry | Add-Member -MemberType NoteProperty -Name AssignmentType -Value "Pool"
        
        $vmEntry
    }
}



function Get-Assignment([string[]]$domainNames)
{
    $vmList = @(Get-WmiObject -Namespace root\cimv2 -Query "Select * from Win32_SessionBrokerTarget where pluginname='VmResource'")
    $vmUserDict = Get-VMAssignmentDictionary $domainNames
    
    $poolDict = @{}
    Get-WmiObject -Namespace root\cimv2\terminalservices -Class Win32_TSRemoteDesktop | %{$poolDict[$_.Alias] = $_.Name}
    
    
    foreach ($vm in $vmList)
    {
        $assignedTo = $NULL
        $assignedType = $NULL

        if ($vmUserDict[$vm.TargetName])
        {
            $assignedTo = $vmUserDict[$vm.TargetName]
            $assignedType = "User"
            
            if ($vm.FarmName -ne "")
            {
                Write-Warning "The virtual machine '$($vm.TargetName)' is assigned to user '$assignedTo' and to pool '$($poolDict[$vm.FarmName])'."
            }
        }
        elseif ($vm.FarmName -ne "")
        {
            $assignedTo = $poolDict[$vm.FarmName]
            $assignedType = "Pool"
        }
    
        $vmEntry = New-Object PSObject

        $vmEntry | Add-Member -MemberType NoteProperty -Name VM -Value $vm.TargetName
        $vmEntry | Add-Member -MemberType NoteProperty -Name Host -Value $vm.Environment
        $vmEntry | Add-Member -MemberType NoteProperty -Name AssignedTo -Value $assignedTo
        $vmEntry | Add-Member -MemberType NoteProperty -Name AssignmentType -Value $assignedType
        
        $vmEntry
    }
}


function Get-UserAssignment
{

param (

    [Parameter(Mandatory=$TRUE)]
    [string[]]
    $User,
    
    [Parameter(Mandatory=$FALSE)]
    [Switch]
    $UserAssigned,
    
    [Parameter(Mandatory=$FALSE)]
    [Switch]
    $UnAssigned

)

    $showAssigned = $UserAssigned -OR -not $UnAssigned
    $showUnAssigned = $UnAssigned -OR -not $UserAssigned

    $users = @($User | %{Get-User $_})
    $domains = $users | %{$_ -split "\\" | Select-Object -First 1} | Sort-Object -Unique
    if ($domains -eq $NULL) { return }

    $vmAssignment = Get-Assignment $domains

    $userAssignment = @{}
    $vmAssignment | ?{$_.AssignmentType -eq "User"} | %{$userAssignment[$_.AssignedTo] = $_}
    
    foreach($entry in $users)
    {
        $vmName = $NULL
        $vmHost = $NULL
        
        if (($showAssigned) -AND ($userAssignment[$entry] -ne $NULL))
        {
            $vmName = ($userAssignment[$entry]).VM
            $vmHost = ($userAssignment[$entry]).Host
        }
        elseif ((-not $showUnAssigned) -OR ($userAssignment[$entry] -ne $NULL))
        {
            continue
        }
        
        $userEntry = New-Object PSObject
        $userEntry | Add-Member -MemberType NoteProperty -Name User -Value $entry
        $userEntry | Add-Member -MemberType NoteProperty -Name VM -Value $vmName
        $userEntry | Add-Member -MemberType NoteProperty -Name Host -Value $vmHost
        
        $userEntry
    }
}


$cbServiceStatus = (Get-Service tssdis -EA SilentlyContinue).Status
if ($cbServiceStatus -ne [System.ServiceProcess.ServiceControllerStatus]::Running)
{
    Write-Error "Connection broker role is not installed or the service is not running."
    exit 1
}

switch ($PsCmdlet.ParameterSetName)
{
    "Users"         {   
                        Get-UserAssignment $User -UserAssigned:$UserAssigned -UnAssigned:$UnAssigned
                        break
                    }

    "Pools"         {
                        $poolNames = Get-WmiObject -Namespace root\cimv2\terminalservices -Class Win32_TSRemoteDesktop | %{$_.Name}
                        $Pool | ?{$poolNames -notcontains $_} | %{Write-Error "Specified pool '$_' does not exist!"}
                        Get-PoolAssignment | ?{$Pool -contains $_.AssignedTo}
                        break
                    }

    "virtual machines"           {
                        $vmNames = Get-WmiObject -Namespace root\cimv2 -Query "Select * from Win32_SessionBrokerTarget where pluginname='VmResource'" | %{$_.TargetName}
                        $VirtualMachine | ?{$vmNames -notcontains $_} | %{Write-Error "Specified virtual machine '$_' does not exist!"}
                        Get-Assignment | ?{$VirtualMachine -contains $_.VM}
                        break
                    }

    "SimpleFilter"  {
                        if (-not $UserAssigned -AND -not $PoolAssigned -AND -not $UnAssigned)
                        {
                            $summary = @{User=0; Pool=0; Free=0}
                            Get-Assignment | %{ if ([string]::IsNullorEmpty($_.AssignmentType)) {$summary['Free']++} else {$summary[$_.AssignmentType]++}; $_}
                            Write-Host "`n`nSummary`n-------"
                            Write-Host "Virtual machines assigned to users : $($summary['User'])"
                            Write-Host "Virtual machines assigned to pools : $($summary['Pool'])"
                            Write-Host "Unassigned Virtual machines        : $($summary['Free'])"
                            Write-Host "Total                 : $(($summary.Values | Measure-Object -Sum).Sum)"
                        }
                        elseif (-not $UserAssigned -AND $PoolAssigned -AND -not $UnAssigned)
                        {
                            Get-PoolAssignment
                        }
                        else
                        {
                            Get-Assignment | ?{($UserAssigned -AND ($_.AssignmentType -eq "User")) `
                                                -OR ($PoolAssigned -AND ($_.AssignmentType -eq "Pool")) `
                                                -OR ($UnAssigned -AND ($_.AssignmentType -eq $NULL))}
                        }
                        break
                    }
}