Scripts/SyncNFSv3ExtendedGroupsFromLDAP.ps1

<#
    .SYNOPSIS
    Synchronizes NFSv3 extended groups from LDAP server.
 
    .DESCRIPTION
    This command syncs NFSv3 extended groups with the user and group data fetched from an LDAP server.
 
    .PARAMETER ResourceGroupName
    The name of the Resource Group in Azure where the storage account resides.
 
    .PARAMETER StorageAccountName
    The name of the Azure storage account.
 
    .PARAMETER LDAPServerIP
    The IP address of the LDAP server.
 
    .PARAMETER LDAPSearchBase
    The search base for querying LDAPServer.
 
    .PARAMETER UserFilter
    Optional filter for selecting users in LDAP server.
 
    .PARAMETER GroupFilter
    Optional filter for selecting groups in LDAP server.
 
    .PARAMETER UidFilePath
    Optional file path for specifying UIDs of users which needs to be synced.
    The file should contain one Uid per line.
 
    .PARAMETER LDAPUserName
    Optional username for authenticating with LDAP server.
    Example: "cn=admin,dc=mydomain,dc=com"
 
    .PARAMETER LDAPCredential
    Optional PSCredential object for authenticating with LDAP server.
 
    .EXAMPLE
    Sync-NFSv3ExtendedGroupsFromLDAP -ResourceGroupName "MyRG" -StorageAccountName "MyStorage" -LDAPServerIP "192.168.1.1" -LDAPSearchBase "dc=mydomain,dc=com"
 
    .NOTES
    Written by: [Azure Blob NFS]
    Date: [October 10, 2024]
#>

function Sync-NFSv3ExtendedGroupsFromLDAP {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory=$true)]
        [string]$ResourceGroupName,

        [Parameter(Mandatory=$true)]
        [string]$StorageAccountName,

        [Parameter(Mandatory=$true)]
        [string]$LDAPServerIP,

        [Parameter(Mandatory=$true)]
        [string]$LDAPSearchBase,

        [Parameter(Mandatory=$false)]
        [string]$UserFilter = "(&(objectClass=posixAccount)(uid=*))",
        
        [Parameter(Mandatory=$false)]
        [string]$GroupFilter = "(&(objectClass=posixGroup)(gidNumber=*))",
        
        [Parameter(Mandatory=$false)]
        [string]$UidFilePath,

        [Parameter(Mandatory=$false, ParameterSetName = 'UserNameSet')]
        [string]$LDAPUserName,

        [Parameter(Mandatory=$false, ParameterSetName = 'PsCredSet')]
        [PSCredential]$LDAPCredential
    )

    # logMessage is written in this way to keep the indentation correct in log file.
    $logMessage = @"
Sync-NFSv3ExtendedGroupsFromLDAP started,
        ResourceGroupName : $ResourceGroupName,
        StorageAccountName : $StorageAccountName,
        LDAPServerIP : $LDAPServerIP,
        LDAPSearchBase : $LDAPSearchBase,
        UserFilter : $UserFilter,
        GroupFilter : $GroupFilter
"@


    if ($UidFilePath)        { $logMessage += ",`n UidFilePath : $UidFilePath" }
    if ($LDAPUserName)        { $logMessage += ",`n LDAPUserName : $LDAPUserName" }
    if ($LDAPCredential)    { $logMessage += ",`n LDAPUserName from PSCredential: $($LDAPCredential.UserName)" }

    Write-NFSv3ExtendedGroupsLog $logMessage -VerbosePreference $VerbosePreference
    
    if ($null -eq $LDAPCredential) {
        # Get admin username and password of LDAP server.
        if ($null -eq $LDAPUserName)
        {
            $LDAPUserName = Read-Host "Enter User Name for LDAP Server Administrator"
        }
        $ldapPassword = Read-Host "Password" -AsSecureString

        # Create a PSCredential object.
        $LDAPCredential = New-Object System.Management.Automation.PSCredential($LDAPUserName, $ldapPassword)
    }

    # Connect to LDAP server.
    try {
        $result = Connect-LdapServer -Server $LDAPServerIP -Credential $LDAPCredential -ErrorAction Stop -Verbose 4>&1 | Out-String
        Write-NFSv3ExtendedGroupsLog "Successfully conneccted to LDAP Server, $result" -VerbosePreference $VerbosePreference 
    } catch {
        # Kill the script if failed to connect to LDAP server.
        Write-NFSv3ExtendedGroupsLog "Connect-LdapServer failed with $_" -LogLevel "ERROR"
        return
    }

    # Create and send LDAP search request for fetching users.
    try{
        $userSearchRequest = New-Object DirectoryServices.Protocols.SearchRequest(
                                $LDAPSearchBase,
                                $UserFilter,
                                [System.DirectoryServices.Protocols.SearchScope]::Subtree)
        $userSearchResult = $LdapConnection.SendRequest($userSearchRequest)

        Write-NFSv3ExtendedGroupsLog "Successfully fetched $($userSearchResult.Entries.Count) user entries from LDAP server" -VerbosePreference $VerbosePreference
    } catch {
        # Kill the script if fetching user entries from LDAP server fails.
        Write-NFSv3ExtendedGroupsLog "SendRequest for user entries failed with error: $_" -LogLevel "ERROR"
        return 
    }

    # Create and send LDAP search request for fetching groups.
    try {
        $groupSearchRequest = New-Object DirectoryServices.Protocols.SearchRequest(
                                $LDAPSearchBase,
                                $groupSearchFilter,
                                [System.DirectoryServices.Protocols.SearchScope]::Subtree)
        $groupSearchResult = $LdapConnection.SendRequest($groupSearchRequest)

        Write-NFSv3ExtendedGroupsLog "Successfully fetched $($groupSearchResult.Entries.Count) group entries from LDAP server" -VerbosePreference $VerbosePreference
    } catch {
        # Kill the script if fetching group entries from LDAP server fails.
        Write-NFSv3ExtendedGroupsLog "SendRequest for group entries failed with error: $_" -LogLevel "ERROR"
        return
    }

    #
    # Create uidToUidNumberMap and uidNumberToExtendedGroupsMap which maps uid to corresponding uidNumber and gidNumber.
    # 1. uidToUidNumberMap maps uid of ldapUser to uidNumber.
    # 2. uidNumberToExtendedGroupsMap maps uidNumber of ldapUser to all the gids it is part of (including primary gid number).
    #
    # For ldapUser1 with uidNumber 1000, primaryGidNumber 1001 and other gids 1002, 1003 and 1004:
    # uidToUidNumberMap[ldapUser1] = 1000
    # uidNumberToExtendedGroupsMap[1000] = [1001, 1002, 1003, 1004]
    #
    $ldapUidNumbers = @()
    $uidToUidNumberMap = @{}
    $uidNumberToExtendedGroupsMap = @{}
    foreach ($entry in $userSearchResult.Entries) {
        if (($null -eq $entry.Attributes["uidNumber"]) -or ($null -eq $entry.Attributes["uid"])) {
            Write-NFSv3ExtendedGroupsLog "Uid or UidNumber is null for $($entry.DistinguishedName)" -LogLevel "WARNING"
            continue
        }
        
        $uidNumber = $entry.Attributes["uidNumber"].GetValues([string])[0]
        $uid = $entry.Attributes["uid"].GetValues([string])[0]

        $uidToUidNumberMap[$uid] = $uidNumber
        $ldapUidNumbers += $uidNumber

        # Add primaryGidNumber of this user to uidNumberToExtendedGroupsMap.
        if (-not ($null -eq $entry.Attributes["gidNumber"])) {
            $primaryGidNumber = $entry.Attributes["gidNumber"].GetValues([string])[0]

            if (-not $uidNumberToExtendedGroupsMap.ContainsKey($uidNumber)) {
                $uidNumberToExtendedGroupsMap[$uidNumber] = @()
            }

            $uidNumberToExtendedGroupsMap[$uidNumber] += $primaryGidNumber
        }
    }

    # Check all the groups each user is part of and store it in uidNumberToExtendedGroupsMap.
    foreach ($entry in $groupSearchResult.Entries) {
        if (($null -eq $entry.Attributes["memberuid"]) -or ($null -eq $entry.Attributes["gidNumber"])) {
            Write-NFSv3ExtendedGroupsLog "GidNumber or MemberUid is null for $($entry.DistinguishedName)" -LogLevel "WARNING"
            continue
        }
        
        # Get gidNumber and all the memberUids which are part of this group.
        $gidNumber = $entry.Attributes["gidnumber"].GetValues([string])[0]
        $memberUidArray = $entry.Attributes["memberuid"].GetValues([string])

        foreach ($memberUid in $memberUidArray) {
            if ($uidToUidNumberMap.ContainsKey($memberUid)) {
                $uidNumber = $uidToUidNumberMap[$memberUid]

                if (-not $uidNumberToExtendedGroupsMap.ContainsKey($uidNumber)) {
                    $uidNumberToExtendedGroupsMap[$uidNumber] = @()
                }

                $uidNumberToExtendedGroupsMap[$uidNumber] += $gidNumber
            }
        }
    }

    # Get all the existing LocalUsers for this account.
    try {
        $localUsers = Get-AzStorageLocalUser -ResourceGroupName $ResourceGroupName -StorageAccountName $StorageAccountName -IncludeNFSv3 -ErrorAction Stop
        Write-NFSv3ExtendedGroupsLog "Fetched $($localUsers.Count) LocalUsers from Azure Storage server" -VerbosePreference $VerbosePreference 
    } catch {
        Write-NFSv3ExtendedGroupsLog "Get-AzStorageLocalUser failed with error: $_" -LogLevel "ERROR"
    }

    # Initialize hashtable to store LocalUser object with name as key.
    $localUserMap = @{}

    # Populate localUserMap hashtable.
    foreach ( $user in $localUsers ) {
        $localUserMap[$user.Name] = $user
    }

    # Get the uids of localusers which needs to be updated from provided file.
    [int[]]$uidsFromFile = @()
    
    if (-not [string]::IsNullOrEmpty($UidFilePath)) {
        $uidsFromFile = Get-Content -Path $UidFilePath | ForEach-Object {
            try {
                [int]$_ 
            }
            catch {
                Write-NFSv3ExtendedGroupsLog -Message "Invalid Uid entry [$_] in $UidFilePath" -LogLevel "ERROR"
            }
        } | Where-Object { $_ -is [int] } # Ensure only integers are kept.
    }

    $logScript = "$PSScriptRoot\WriteNFSv3ExtendedGroupsLog.ps1"
    $jobList = [System.Collections.Generic.List[System.Management.Automation.Job]]::new()
    
    #
    # There is a limit of 10 Write Ops per second by using ARM with Azure Storage. To avoid requests getting
    # throttled because of this limit, keep maxParallelJobs less than or equal to 10.
    #
    # https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/request-limits-and-throttling#storage-throttling
    #
    $maxParallelJobs = 10
    $ldapUidNumbersSize = $ldapUidNumbers.Length
    $chunkUidNumbersSize = [Math]::Ceiling($ldapUidNumbersSize / $maxParallelJobs)
    $numOfIterations = [Math]::Ceiling($ldapUidNumbersSize / $chunkUidNumbersSize)

    for ( $i = 0; $i -lt $numOfIterations; $i++ ) {
        # Slice the ldapUidNumbers array into equal size of chunks each of which will be handled by separate processes.
        $startIndex = $i * $chunkUidNumbersSize
        $endIndex = [Math]::Min($startIndex + $chunkUidNumbersSize - 1, $ldapUidNumbersSize - 1)

        $chunkUidNumbers = $ldapUidNumbers[$startIndex..$endIndex]
        
        # Start a new background job for each chunk of users.
        $job = Start-Job -ScriptBlock {
            param (
                $chunkUidNumbers,
                $uidNumberToExtendedGroupsMap,
                $localUserMap,
                $ResourceGroupName,
                $StorageAccountName,
                $logScript,
                $uidsFromFile,
                $VerbosePreference
            )
            
            # Dot-Source internal logging function and import required modules cause this background job will create different runspace.
            Import-Module Az.Storage 4>&1 | Out-Null
            $Global:NFSv3ExtendedGroupsLogFile = $using:NFSv3ExtendedGroupsLogFile
            . $logScript

            #
            # Convert uidsFromFile to HashMap for O(1) search. Can not pass complex HashMap objects to background
            # jobs, therefore it was initially stored in an array.
            #
            $uidsFromFileSet = [System.Collections.Generic.HashSet[int]]::new()
            foreach ( $uid in $uidsFromFile ) {
                $uidsFromFileSet.Add($uid) | Out-Null
            }

            foreach ($uidNumber in $chunkUidNumbers) {
                #
                # If user has provided file to tell which Uids to upload,
                # Verify this Uid is present in that list and skip if not.
                #
                if ( $uidsFromFileSet.Count -gt 0 ) {
                    if ( -not $uidsFromFileSet.Contains($uidNumber) ) {
                        # If this uid is not part of list provided by user, skip.
                        Write-NFSv3ExtendedGroupsLog "$uidNumber is not present in the file, skipping" -VerbosePreference $VerbosePreference
                        continue
                    }
                }
                
                # Get extended group gids for this uid.
                $gidArray = @()
                if ($uidNumberToExtendedGroupsMap.ContainsKey($uidNumber)) {
                    $gidArray = $uidNumberToExtendedGroupsMap[$uidNumber]
                }

                # Don't process if extended groups array is empty.
                if (-not $gidArray -or $gidArray.Count -eq 0) {
                    Write-NFSv3ExtendedGroupsLog "ExtendedGroups array is null or empty for Uid $uidNumber" -VerbosePreference $VerbosePreference
                    continue
                }

                # Convert ArrayList to Int32 array.
                [int32[]]$extendedGroupsArray = [int32[]]$gidArray | Sort-Object
                
                # NFSv3 localuser have name in "nfsv3_<uid>" format.
                $localUserName = "nfsv3_$($uidNumber)"

                # Check if local user with same uidNumber exist in localuser database.
                if ( $localUserMap.ContainsKey($localUserName) ) {
                    # Sort extended group array before comparing.
                    $existingLocalUser = $localUserMap[$localUserName]
                    $existingExtendedGroups = $existingLocalUser.ExtendedGroups | Sort-Object
                    $difference = Compare-Object -ReferenceObject $existingExtendedGroups -DifferenceObject $extendedGroupsArray

                    # Don't PUT local user if database is already updated.
                    if ($difference.Count -eq 0) {
                        $extendedGroupsArrayCommaSeparated = $extendedGroupsArray -join ','
                        Write-NFSv3ExtendedGroupsLog "ExtendedGroups already up to date for Uid $uidNumber, ExtendedGroups: [$extendedGroupsArrayCommaSeparated]" -VerbosePreference $VerbosePreference
                        continue
                    }
                } 

                # Create/update localuser.
                try {
                    Set-AzStorageLocalUser -ResourceGroupName $ResourceGroupName -StorageAccountName $StorageAccountName -UserName $localUserName -IsNfSv3Enabled $true -ExtendedGroups $extendedGroupsArray -ErrorAction Stop | Out-Null
                    
                    # On success.
                    $extendedGroupsArrayCommaSeparated = $extendedGroupsArray -join ','
                    Write-NFSv3ExtendedGroupsLog "Set-AzStorageLocalUser completed successfully for /$ResourceGroupName/$StorageAccountName/$localUserName, ExtendedGroups: [$extendedGroupsArrayCommaSeparated]" -VerbosePreference $VerbosePreference
                } catch {
                    # On failure.
                    Write-NFSv3ExtendedGroupsLog "Set-AzStorageLocalUser failed for /$ResourceGroupName/$StorageAccountName/$localUserName with error: $_" -LogLevel "ERROR"
                }
            }
        } -ArgumentList $chunkUidNumbers, $uidNumberToExtendedGroupsMap, $localUserMap, $ResourceGroupName, $StorageAccountName, $logScript, $uidsFromFile, $VerbosePreference
    
        $jobList.Add($job)
    }

    # Wait for all jobs to complete.
    Write-NFSv3ExtendedGroupsLog "Waiting for all jobs to complete..." -VerbosePreference $VerbosePreference
    $jobList | ForEach-Object { $_ | Wait-Job } | Out-Null
    Write-NFSv3ExtendedGroupsLog "All jobs completed" -VerbosePreference $VerbosePreference

    # Receive the job output and write it to the console.
    $jobList | ForEach-Object { Receive-Job -Job $_ }

    # Clean up completed jobs.
    $jobList | ForEach-Object { $_ | Remove-Job }

    Write-NFSv3ExtendedGroupsLog "Sync-NFSv3ExtendedGroupsFromLDAP completed" -VerbosePreference $VerbosePreference
}