Framework/Managers/ControlStateExtension.ps1

Set-StrictMode -Version Latest

class ControlStateExtension
{
    hidden [PSObject] $AzSDKResourceGroup = $null;
    hidden [PSObject] $AzSDKStorageAccount = $null;
    hidden [PSObject] $AzSDKStorageContainer = $null;
    hidden [PSObject] $ControlStateIndexer = $null;    
    hidden [int] $HasControlStateReadPermissions = -1;
    hidden [int] $HasControlStateWritePermissions = -1;
    hidden [string]    $IndexerBlobName ="Resource.index.json"


    ControlStateExtension()
    {

    }

    hidden [void] Initialize()
    {
        try
        {
            $this.GetAzSDKControlStateContainer()
            $this.HasControlStateReadPermissions = 1

            #If there are no write permissions and with just read permissions, we cant get the container object.
            if($null -ne $this.AzSDKResourceGroup -and $null -ne $this.AzSDKStorageAccount -and $null -ne $this.AzSDKStorageContainer)
            {
                $this.HasControlStateWritePermissions = 1
            }
        }
        catch
        {
            $this.HasControlStateReadPermissions = 0
        }
    }

    hidden [PSObject] GetAzSDKRG()
    {
        $azSDKConfigData = [ConfigurationManager]::GetAzSdkConfigData()

        $resourceGroup = Get-AzureRmResourceGroup -Name $azSDKConfigData.AzSDKRGName -ErrorAction SilentlyContinue
        if($null -eq $resourceGroup -or ($resourceGroup | Measure-Object).Count -eq 0)
        {
            if([Helpers]::NewAzSDKResourceGroup($azSDKConfigData.AzSDKRGName, [Constants]::AzSDKRGLocation, ""))
            {
                $resourceGroup = Get-AzureRmResourceGroup -Name $azSDKConfigData.AzSDKRGName -ErrorAction SilentlyContinue
            }
        }
        $this.AzSDKResourceGroup = $resourceGroup
        return $resourceGroup;
    }

    hidden [PSObject] GetAzSDKStorageAccount()
    {
        if($null -eq $this.AzSDKResourceGroup)
        {
            $this.GetAzSDKRG();
        }
        $StorageAccount = Get-AzureRmStorageAccount -ResourceGroupName $this.AzSDKResourceGroup.ResourceGroupName | Where-Object {$_.StorageAccountName -like 'azsdk*'} -ErrorAction SilentlyContinue 
        #if no storage account found then it assumes that there is no control state feature is not used and if there are more than one storage account found it assumes the same
        if($null -eq $StorageAccount -or ($StorageAccount | Measure-Object).Count -eq 0)
        {
            $storageAccountName = ("azsdk" + (Get-Date).ToUniversalTime().ToString("yyyyMMddHHmmss"));    
            $storageObject = [Helpers]::NewAzsdkCompliantStorage($storageAccountName, $this.AzSDKResourceGroup.ResourceGroupName, [Constants]::AzSDKRGLocation)
            if($null -ne $storageObject -and ($storageObject | Measure-Object).Count -gt 0)
            {
                $StorageAccount = Get-AzureRmStorageAccount -ResourceGroupName $this.AzSDKResourceGroup.ResourceGroupName | Where-Object {$_.StorageAccountName -like 'azsdk*'} -ErrorAction SilentlyContinue                     
            }                    
        }
        $this.AzSDKStorageAccount = $StorageAccount;
        return $StorageAccount;
    }

    hidden [PSObject] GetAzSDKControlStateContainer()
    {
        if($null -eq $this.AzSDKStorageAccount)
        {
            $this.GetAzSDKStorageAccount();
        }

        $ContainerName = "azsdk-controls-state"
        $containerObject = Get-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction SilentlyContinue
        if($null -eq $containerObject)
        {
            New-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction SilentlyContinue
            $containerObject = $containerObject = Get-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction SilentlyContinue
        }
        $this.AzSDKStorageContainer = $containerObject;
        return $containerObject;
    }
    
    hidden [bool] ComputeControlStateIndexer()
    {
        #check for permission validaiton

        if($this.HasControlStateReadPermissions -le 0) 
        {
            return $false;
        }
                    
        $StorageAccount = $this.AzSDKStorageAccount;
        $containerObject = $this.AzSDKStorageContainer;
        $ContainerName = $this.AzSDKStorageContainer.Name
        $indexerBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $this.IndexerBlobName -Context $StorageAccount.Context -ErrorAction SilentlyContinue
        [ControlStateIndexer[]] $indexerObjects = @();
        $this.ControlStateIndexer  = $indexerObjects
        if($null -eq $indexerBlob)
        {            
            return $true;
        }
        $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";
        if(-not (Test-Path -Path $AzSDKTemp))
        {
            mkdir -Path $AzSDKTemp -Force
        }
        Get-AzureStorageBlobContent -CloudBlob $indexerBlob.ICloudBlob -Context $StorageAccount.Context -Destination $AzSDKTemp -Force
        $indexerObject = Get-ChildItem -Path "$AzSDKTemp\$($this.IndexerBlobName)" -Force | Get-Content | ConvertFrom-Json 
        $this.ControlStateIndexer += $indexerObject;
        return $true;
    }

    hidden [PSObject] GetControlState([string] $id)
    {
        try
        {
            [ControlState[]] $controlStates = @();
            $retVal = $this.ComputeControlStateIndexer();
            if($null -ne $this.ControlStateIndexer -and  $retVal)
            {
                $indexes = @();
                $indexes += $this.ControlStateIndexer 

                $selectedIndex = $indexes | Where-Object { $_.ResourceId -eq $id}

                if(($selectedIndex | Measure-Object).Count -gt 0)
                {
                    $controlStateBlobName = $selectedIndex.HashId + ".json"
                    $azSDKConfigData = [ConfigurationManager]::GetAzSdkConfigData()
                    #$azSDKConfigData.$AzSDKRGName
                    #Look of is there is a AzSDK RG and AzSDK Storage account
                    $StorageAccount = $this.AzSDKStorageAccount;                        
                    $containerObject = $this.AzSDKStorageContainer
                    $ContainerName = $this.AzSDKStorageContainer.Name
                    $controlStateBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $controlStateBlobName -Context $StorageAccount.Context -ErrorAction SilentlyContinue

                    if($null -eq $controlStateBlob)
                    {
                        return $controlStates;
                    }
                    $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";
                    if(-not (Test-Path -Path $AzSDKTemp))
                    {
                        mkdir -Path $AzSDKTemp -Force
                    }
                    Get-AzureStorageBlobContent -CloudBlob $controlStateBlob.ICloudBlob -Context $StorageAccount.Context -Destination $AzSDKTemp -Force
                    $ControlStatesJson = Get-ChildItem -Path "$AzSDKTemp\$controlStateBlobName" -Force | Get-Content | ConvertFrom-Json 

                    if($null -ne $ControlStatesJson)
                    {                    
                        $ControlStatesJson | ForEach-Object {
                            $controlStates += [ControlState] $_;
                        }
                    }
                }
            }
            return $controlStates;
        }
        finally{
            $this.CleanTempFolder();
        }
    }

    hidden [void] SetControlState([string] $id, [ControlState[]] $controlStates, [bool] $Override)
    {        
        $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";                
        if(-not (Test-Path "$AzSDKTemp\ControlState"))
        {
            mkdir -Path "$AzSDKTemp\ControlState" -ErrorAction Stop | Out-Null
        }

        $hash = [Helpers]::ComputeHash($id);
        $indexerPath = "$AzSDKTemp\ControlState\$($this.IndexerBlobName)"
        $fileName = "$AzSDKTemp\ControlState\$hash.json"    
        
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = $this.AzSDKStorageContainer.Name
        $finalControlStates = $controlStates;
        if($Override)
        {
            # in the case of override, just persist what is evaluated in the current context. No merging with older data
            $this.UpdateControlIndexer($id, $controlStates);
            $finalControlStates = $controlStates | Where-Object { $_.State};
        }
        else
        {
            #merge with the exiting if found
            $persistedControlStates = $this.GetPersistedControlStates("$hash.json");
            $finalControlStates = $this.MergeControlStates($persistedControlStates, $controlStates);
            $this.UpdateControlIndexer($id, $finalControlStates);
        }
        [Helpers]::ConvertToJsonCustom($finalControlStates) | Out-File $fileName -Force        

        if($null -ne $this.ControlStateIndexer)
        {                
            [Helpers]::ConvertToJsonCustom($this.ControlStateIndexer) | Out-File $indexerPath -Force
            $controlStateArray = Get-ChildItem -Path "$AzSDKTemp\ControlState"                
            $controlStateArray | ForEach-Object { Set-AzureStorageBlobContent -File $_.FullName -Container $ContainerName -BlobType Block -Context $StorageAccount.Context -Force }
        }
        else
        {
            #clean up the container as there is no indexer
            Get-AzureStorageBlob -Container $ContainerName -Context $StorageAccount.Context | Remove-AzureStorageBlob  
        }
    }

    hidden [ControlState[]] GetPersistedControlStates([string] $controlStateBlobName)
    {
        $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";
        if(-not (Test-Path "$AzSDKTemp\ExistingControlStates"))
        {
            mkdir -Path "$AzSDKTemp\ExistingControlStates" -ErrorAction Stop | Out-Null
        }
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = $this.AzSDKStorageContainer.Name
        [ControlState[]] $ControlStatesJson = @()
        try
        {
            $controlStateBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $controlStateBlobName -Context $StorageAccount.Context -ErrorAction SilentlyContinue
            Get-AzureStorageBlobContent -CloudBlob $controlStateBlob.ICloudBlob -Context $StorageAccount.Context -Destination "$AzSDKTemp\ExistingControlStates" -Force
            $ControlStatesJson = Get-ChildItem -Path "$AzSDKTemp\ExistingControlStates\$controlStateBlobName" -Force | Get-Content | ConvertFrom-Json 
        }
        catch
        {
            $ControlStatesJson = @()
        }
        return $ControlStatesJson
    }

    hidden [ControlState[]] MergeControlStates([ControlState[]] $persistedControlStates,[ControlState[]] $controlStates)
    {
        [ControlState[]] $computedControlStates = $controlStates;
        if(($computedControlStates | Measure-Object).Count -le 0)
        {
            $computedControlStates = @();
        }
        if(($persistedControlStates | Measure-Object).Count -gt 0)
        {
            $persistedControlStates | ForEach-Object {
                $controlState = $_;
                if(($computedControlStates | Where-Object { $_.InternalId -eq $controlState.InternalId} | Measure-Object).Count -le 0)
                {
                    $computedControlStates += $controlState;
                }
            }
        }
        #remove the control states with null state which would be in the case of clear attestation.
        $computedControlStates = $computedControlStates | Where-Object { $_.State}

        return $computedControlStates;
    }

    hidden [void] UpdateControlIndexer([string] $id, [ControlState[]] $controlStates)
    {
        $this.ControlStateIndexer = $null;
        $retVal = $this.ComputeControlStateIndexer();
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = $this.AzSDKStorageContainer.Name
        if($retVal)
        {                
            $tempHash = [Helpers]::ComputeHash($id);
            $filteredIndexerObject = $this.ControlStateIndexer | Where-Object { $_.HashId -eq $tempHash}
            if($null -ne $filteredIndexerObject)
            {
                if(($controlStates | Measure-Object).Count -le 0)
                {
                    $this.ControlStateIndexer = $this.ControlStateIndexer | Where-Object { $_.HashId -ne $tempHash}
                }
                else
                {
                    $filteredIndexerObject.ExpiryTime = [DateTime]::UtcNow.AddMonths(3);
                    $filteredIndexerObject.AttestedBy =  [Helpers]::GetCurrentSessionUser();
                    $filteredIndexerObject.AttestedDate = [DateTime]::UtcNow;
                    $filteredIndexerObject.Version = "1.0";
                }
            }
            else
            {
                $currentIndexObject = [ControlStateIndexer]::new();
                $currentIndexObject.ResourceId = $id
                $currentIndexObject.HashId = $tempHash;
                $currentIndexObject.ExpiryTime = [DateTime]::UtcNow.AddMonths(3);
                $currentIndexObject.AttestedBy = [Helpers]::GetCurrentSessionUser();
                $currentIndexObject.AttestedDate = [DateTime]::UtcNow;
                $currentIndexObject.Version = "1.0";
                $this.ControlStateIndexer += $currentIndexObject;            
            }
        }
    }
    
    [bool] HasControlStateReadAccessPermissions()
    {
        if($this.HasControlStateReadPermissions -le 0)
        {
            return $false;
        }
        else
        {
            return $true;
        }
    }

    [bool] HasControlStateWriteAccessPermissions()
    {        
        if($this.HasControlStateWritePermissions -le 0)
        {
            return $false;
        }
        else
        {
            return $true;
        }
    }

    hidden [void] CleanTempFolder()
    {
        $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp";                
        if(Test-Path "$AzSDKTemp")
        {
            rmdir -Path $AzSDKTemp -Recurse -Force -ErrorAction Stop | Out-Null
        }

    }

    
}