Private/datalake-acls.ps1

using namespace Azure.Storage.Files.DataLake.Models

$script:Global_ctxStorageAccountName = ""
$script:Global_ctxContainerName = ""
$script:Global_ctx = ""

function Get-RKStorageContext {
    [CmdletBinding()]
    param([string]$StorageAccountName,
        [string]$ContainerName)

    #Wrap call in $Null to prevent any additional piped output from being sent back as context.
    $Null = @(
        # If the storage account in use has changed update the Context
        if (($script:Global_ctxStorageAccountName -ne $StorageAccountName) `
                -or ($script:Global_ctxContainerName -ne $ContainerName)) {
            # Flush the cache if changing
            Update-RKLakeFromCache
            $script:Global_ctx = New-AzStorageContext -StorageAccountName $StorageAccountName -UseConnectedAccount
            $script:Global_ctxStorageAccountName = $StorageAccountName
            $script:Global_ctxContainerName = $ContainerName
            Reset-RKCache
        }
    )
    return $script:Global_ctx
}




function Get-RKLakePathAcl {
    [CmdletBinding()]
    param (
        [string]$Path,
        [string]$ContainerName,
        $ctx
    )
    
    if ("" -eq $Path) {
        $Path = '/'
    }
    if ($script:aclCache.ContainsKey($Path)) {
        return $script:aclCache[$Path]
    }
    elseif ($Path -eq '/') {
        $acl = (Get-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName).ACL
    }
    else {
        $acl = (Get-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName -Path $Path).ACL
    }
    $script:aclCache[$Path] = $acl
    return $acl
}


function Test-RKLakePermission {
    [CmdletBinding()]
    param(
        [string]$Path,
        [string]$AccessControlType,
        [string]$EntityId,
        [RolePermissions]$Permissions,
        [bool]$DefaultScope = $False,
        [string]$ContainerName,
        $ctx
    )
    $Current = Get-RKLakePathAcl -Path $Path -ContainerName $ContainerName -ctx $ctx
    $MatchedAcls = $Current | Where-Object { (($_.AccessControlType -eq $AccessControlType) `
                -and ($_.EntityId -eq $EntityId ) `
                -and ($_.DefaultScope -eq $DefaultScope ) 
        ) }

    if (-not $MatchedAcls) {
        return $false
    }
    else {
        return (($MatchedAcls.Permissions) -eq $Permissions)
    }
    
}

function ConvertTo-RKRolePermissions {
    [CmdletBinding()]
    param($TargetPermissions)

    if ($TargetPermissions.Length -gt 0) {
        $TargetPermissions = $TargetPermissions.replace("r", 'Read,').replace("w", 'Write,').replace("x", 'Execute,')
        $TargetPermissions = $TargetPermissions.Substring(0, $TargetPermissions.Length - 1)
        $TargetPermissions = [RolePermissions]$TargetPermissions
    }
    else {
        $TargetPermissions = [RolePermissions]::None
    }
    return $TargetPermissions
}

function ConvertTo-RKPermissionString {
    [CmdletBinding()]
    param(
        [RolePermissions]$Permissions
    )
    
    if (($Permissions -band 'Read') -eq 'Read') {
        $Res = "r"
    } 
    else {
        $Res = "-"
    }
    if (($Permissions -band 'Write') -eq 'Write') {
        $Res = $Res + "w"
    } 
    else {
        $Res = $Res + "-"
    }
    if (($Permissions -band 'Execute') -eq 'Execute') {
        $Res = $Res + "x"
    } 
    else {
        $Res = $Res + "-"
    }
    return $Res
}

function Get-RKUpperPath {
    [CmdletBinding()]
    param([string]$Path)
    if (-not $Path.Contains('/')) {
        return $null
    }
    $Paths = $path.split('/')
    return $Paths[0..($Paths.Count - 2)] -join '/'
}


function Add-RKLakePathAcl {
    [CmdletBinding()]
    param(
        [string]$Path,
        [string]$EntityId,
        [string]$AccessControlType,
        [string]$Permissions,
        [bool]$SkipUpperFolders = $False,
        [bool]$DefaultPermission = $False,
        [bool]$ApplyToSubFolders = $True,
        [string]$ContainerName,
        $ctx
    )
    $Path = $Path.Trim("/")
    $PermissionsObj = ConvertTo-RKRolePermissions $Permissions
    $PermissionsStr = ConvertTo-RKPermissionString $PermissionsObj
    $Exists = Test-RKLakePermission -ContainerName $ContainerName -ctx $ctx -Path $path -AccessControlType $AccessControlType -EntityId $EntityId -Permissions $PermissionsObj -DefaultScope $DefaultPermission
    if (-not $Exists) {
        if ($ApplyToSubFolders) {
            if ($Path -eq "") { $Path = '/' }
            if ($script:recursiveAclsToApply.ContainsKey($Path)) {
                $acl = $script:recursiveAclsToApply[$Path]
                $acl = Set-AzDataLakeGen2ItemAclObject -AccessControlType $AccessControlType -Permission $PermissionsStr -EntityId $EntityId -DefaultScope:$DefaultPermission -InputObject $acl
            }
            else {
                $acl = Set-AzDataLakeGen2ItemAclObject -AccessControlType $AccessControlType -Permission $PermissionsStr -EntityId $EntityId -DefaultScope:$DefaultPermission
            }
            $script:recursiveAclsToApply[$path] = $acl
        }
        else {
            if ($Path -eq "") { $Path = '/' }
            $script:aclCache[$Path] = Set-AzDataLakeGen2ItemAclObject -AccessControlType $AccessControlType -EntityID $EntityId -Permission $PermissionsStr -InputObject $script:aclCache[$path] -DefaultScope:$DefaultPermission
            $script:changedPaths += $Path
        }
    }
    if ((-not $SkipUpperFolders) -and ($Path -ne "")) {
        $Path = Get-RKUpperPath -Path $Path
        Add-RKLakePathAcl -ContainerName $ContainerName -ctx $ctx -Path $Path -EntityId $EntityId -AccessControlType $AccessControlType -Permissions "rx" -ApplyToSubFolders $False -DefaultPermission $False -SkipUpperFolders $False
    }
    
}

function Test-RKFolderExists {
    [CmdletBinding()]
    param(
        $Path,
        [string]$ContainerName,
        $ctx
    )
    if (Get-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName -Path $Path -ErrorAction:SilentlyContinue) {
        return $true
    }
    else {
        return $false
    }
}

function New-RKFolder {
    [CmdletBinding()]
    param(
        $Path,
        [string]$ContainerName,
        $ctx
    )
    $dir = New-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName -Path $path -Directory
}


function Set-RKLakePathAcl {
    [CmdletBinding()]
    param(
        [string]$ObjectType,
        [string]$DisplayName,
        [string]$Path,
        [string]$Permissions,
        [string]$ContainerName,
        $ctx
    )
    $EntityId = Get-RKAADObjectId -ObjectType $ObjectType -DisplayName $DisplayName

    if (-not (Test-RKFolderExists -Path $Path -ctx $ctx -ContainerName $ContainerName)) {
        New-RKFolder -Path $Path -ctx $ctx -ContainerName $ContainerName
    }

    if ($ObjectType -eq "Group") {
        $AccessControlType = "group"
    }
    else {
        $AccessControlType = "user"
    }

    # Apply permissions on existing items (apply default above, and recurse down)
    Add-RKLakePathAcl -Path $Path -EntityId $EntityId -AccessControlType $AccessControlType `
        -Permissions $Permissions -SkipUpperFolders $False -DefaultPermission $False -ApplyToSubFolders $true `
        -ContainerName $ContainerName -ctx $ctx
    # Apply default permissions for new items (no need to apply above, but must recurse)
    Add-RKLakePathAcl -Path $Path -EntityId $EntityId -AccessControlType $AccessControlType `
        -Permissions $Permissions -SkipUpperFolders $true -DefaultPermission $true -ApplyToSubFolders $true `
        -ContainerName $ContainerName -ctx $ctx

}

function Reset-RKCache {
    [CmdletBinding()]
    $script:aclCache = @{}
    $script:recursiveAclsToApply = @{}
    $script:changedPaths = @()
}

function Update-RKLakeFromCache {
    # Parameter help description
    $ContainerName = $script:Global_ctxContainerName
    $ctx = $script:Global_ctx
    
    foreach ($Path in $script:changedPaths) {
        $acl = $script:aclCache[$Path]
        Write-Host "Updating $Path Single ACL"
        if ($Path -eq "/") {
            Update-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName -Acl $acl
        }
        else {
            Update-AzDataLakeGen2Item -Context $ctx -FileSystem $ContainerName -Path $Path -Acl $acl
        }
    }

    foreach ($Path in $script:recursiveAclsToApply.Keys) {
        $acl = $script:recursiveAclsToApply[$Path]
        Write-Host "Updating $Path Recursively"
        if ($Path -eq "/") {
            Update-AzDataLakeGen2AclRecursive  -Context $ctx -FileSystem $ContainerName -Acl $acl
        }
        else {
            Update-AzDataLakeGen2AclRecursive  -Context $ctx -FileSystem $ContainerName -Path $Path -Acl $acl
        }   
    }

    Reset-RKCache
}