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"
    hidden [int] $retryCount = 3;


    ControlStateExtension()
    {

    }

    hidden [void] Initialize([bool] $CreateResourcesIfNotExists)
    {
        $this.GetAzSDKControlStateContainer($CreateResourcesIfNotExists)
    }

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

        $resourceGroup = Get-AzureRmResourceGroup -Name $azSDKConfigData.AzSDKRGName -ErrorAction SilentlyContinue
        if($createIfNotExists -and ($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 [void] GetAzSDKStorageAccount($createIfNotExists)
    {
        if($null -eq $this.AzSDKResourceGroup)
        {
            $this.GetAzSDKRG($createIfNotExists);
        }
        if($null -ne $this.AzSDKResourceGroup)
        {
            $StorageAccount  = $null;
            $loopValue = $this.retryCount;
            while($loopValue -gt 0)
            {
                $loopValue = $loopValue - 1;
                try
                {
                    $StorageAccount = Get-AzureRmStorageAccount -ResourceGroupName $this.AzSDKResourceGroup.ResourceGroupName -ErrorAction Stop | Where-Object {$_.StorageAccountName -like 'azsdk*'} -ErrorAction Stop 
                    $loopValue = 0;
                }
                catch
                {
                    #eat this exception and retry
                }
            }            

            #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($createIfNotExists -and ($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)
                {
                    $loopValue = $this.retryCount;
                    while($loopValue -gt 0)
                    {
                        $loopValue = $loopValue - 1;
                        try
                        {
                            $StorageAccount = Get-AzureRmStorageAccount -ResourceGroupName $this.AzSDKResourceGroup.ResourceGroupName -ErrorAction Stop | Where-Object {$_.StorageAccountName -like 'azsdk*'} -ErrorAction Stop                     
                            $loopValue = 0;
                        }
                        catch
                        {
                            #eat this exception and retry
                        }
                    }                    
                }                    
            }
            $this.AzSDKStorageAccount = $StorageAccount;
        }
    }

    hidden [void] GetAzSDKControlStateContainer([bool] $createIfNotExists)
    {
        $ContainerName = [Constants]::StateContainerName;
        if($null -eq $this.AzSDKStorageAccount)
        {
            $this.GetAzSDKStorageAccount($createIfNotExists)
        }
        if($null -eq $this.AzSDKStorageAccount)
        {
            #No storage account => no permissions at all
            $this.HasControlStateReadPermissions = 0
            $this.HasControlStateWritePermissions = 0
            return;
        }
        
        $this.HasControlStateReadPermissions = 0                    
        $this.HasControlStateWritePermissions = 0
        $writeTestContainerName = "writetest";

        #see if user can create the test container in the storage account. If yes then user have both RW permissions.
        try
        {
            $containerObject = Get-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $writeTestContainerName -ErrorAction SilentlyContinue
            if($null -ne $containerObject)
            {
                Remove-AzureStorageContainer -Name $writeTestContainerName -Context  $this.AzSDKStorageAccount.Context -ErrorAction Stop -Force
                $this.HasControlStateWritePermissions = 1
                $this.HasControlStateReadPermissions = 1
            }
            else
            {
                New-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $writeTestContainerName -ErrorAction Stop
                $this.HasControlStateWritePermissions = 1
                $this.HasControlStateReadPermissions = 1
                Remove-AzureStorageContainer -Name $writeTestContainerName -Context  $this.AzSDKStorageAccount.Context -ErrorAction SilentlyContinue -Force
            }                
        }
        catch
        {
            $this.HasControlStateWritePermissions = 0
        }
        if($this.HasControlStateWritePermissions -eq 1)
        {
            try
            {
                if($createIfNotExists)
                {
                    New-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction SilentlyContinue
                }
                $containerObject = Get-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction SilentlyContinue
                $this.AzSDKStorageContainer = $containerObject;                    
            }
            catch
            {
                #Do nothing
            }
        }
        else
        {
            #if user doesnt have writepermission, check atleast user have read permission
            try
            {
                #Able to read the container then read permissions are good
                $containerObject = Get-AzureStorageContainer -Context $this.AzSDKStorageAccount.Context -Name $ContainerName -ErrorAction Stop
                $this.AzSDKStorageContainer = $containerObject;
                $this.HasControlStateReadPermissions = 1
            }
            catch
            {
                #Resetting permissions in the case of exception
                $this.HasControlStateReadPermissions = 0            
            }    
        }        
    }
    
    hidden [bool] ComputeControlStateIndexer()
    {
        #check for permission validation
        if($this.HasControlStateReadPermissions -le 0) 
        {
            return $false;
        }
        #return if you don't have the required state attestation configuration during the runtime evaluation
        if( $null -eq $this.AzSDKResourceGroup -or $null -eq $this.AzSDKStorageAccount -or $null -eq $this.AzSDKStorageContainer)
        {
            return $false;
        }
        $StorageAccount = $this.AzSDKStorageAccount;
        $containerObject = $this.AzSDKStorageContainer;
        $ContainerName = ""
        if($null -ne $this.AzSDKStorageContainer)
        {
            $ContainerName = $this.AzSDKStorageContainer.Name
        }        
        $indexerBlob = $null;        
        $loopValue = $this.retryCount;
        while($loopValue -gt 0)
        {
            $loopValue = $loopValue - 1;
            try
            {
                $indexerBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $this.IndexerBlobName -Context $StorageAccount.Context -ErrorAction Stop
                $loopValue = 0;
            }
            catch
            {
                #Do Nothing. Below code would create a default indexer.
            }
        }
        [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
        }

        $indexerObject = @();
        $loopValue = $this.retryCount;
        while($loopValue -gt 0)
        {
            $loopValue = $loopValue - 1;
            try
            {
                Get-AzureStorageBlobContent -CloudBlob $indexerBlob.ICloudBlob -Context $StorageAccount.Context -Destination $AzSDKTemp -Force -ErrorAction Stop
                $indexerObject = Get-ChildItem -Path "$AzSDKTemp\$($this.IndexerBlobName)" -Force -ErrorAction Stop | Get-Content | ConvertFrom-Json
                $loopValue = 0;
            }
            catch
            {
                #eat this exception and retry
            }
        }
        $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 -and $_.ExpiryTime -gt [DateTime]::UtcNow}

                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 = ""
                    if($null -ne $this.AzSDKStorageContainer)
                    {
                        $ContainerName = $this.AzSDKStorageContainer.Name
                    }

                    $loopValue = $this.retryCount;
                    $controlStateBlob = $null;
                    while($loopValue -gt 0 -and $controlStateBlob -eq $null)
                    {
                        $loopValue = $loopValue - 1;
                        $controlStateBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $controlStateBlobName -Context $StorageAccount.Context -ErrorAction SilentlyContinue
                    }

                    if($null -eq $controlStateBlob)
                    {
                        return $null;
                    }
                    $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";
                    if(-not (Test-Path -Path $AzSDKTemp))
                    {
                        mkdir -Path $AzSDKTemp -Force
                    }

                    $loopValue = $this.retryCount;
                    while($loopValue -gt 0)
                    {
                        $loopValue = $loopValue - 1;
                        try
                        {
                            Get-AzureStorageBlobContent -CloudBlob $controlStateBlob.ICloudBlob -Context $StorageAccount.Context -Destination $AzSDKTemp -Force -ErrorAction Stop
                            $loopValue = 0;
                        }
                        catch
                        {
                            #eat this exception and retry
                        }
                    }
                    $ControlStatesJson = Get-ChildItem -Path "$AzSDKTemp\$controlStateBlobName" -Force | Get-Content | ConvertFrom-Json 
                    if($null -ne $ControlStatesJson)
                    {                    
                        $ControlStatesJson | ForEach-Object {
                            try
                            {
                                $controlState = [ControlState] $_

                                #this can be removed after we check there is no value for attestationstatus coming to azsdktm database
                                if($controlState.AttestationStatus -eq [AttestationStatus]::NotFixed)
                                {
                                    $controlState.AttestationStatus = [AttestationStatus]::WillNotFix
                                }

                                $controlStates += $controlState;                                
                            }
                            catch 
                            {
                                [EventBase]::PublishGenericException($_);
                            }
                        }
                    }
                }
            }
            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
        }
        else
        {
            Remove-Item -Path "$AzSDKTemp\ControlState\*" -Force -Recurse 
        }

        $hash = [Helpers]::ComputeHash($id);
        $indexerPath = "$AzSDKTemp\ControlState\$($this.IndexerBlobName)"
        $fileName = "$AzSDKTemp\ControlState\$hash.json"    
        
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = ""
        if($null -ne $this.AzSDKStorageContainer)
        {
            $ContainerName = $this.AzSDKStorageContainer.Name
        }


        #Filter out the states where there is
        $finalControlStates = $controlStates | Where-Object { $_.ActualVerificationResult -ne [VerificationResult]::Passed};
        if(($finalControlStates | Measure-Object).Count -gt 0)
        {
            if($Override)
            {
                # in the case of override, just persist what is evaluated in the current context. No merging with older data
                $this.UpdateControlIndexer($id, $finalControlStates, $false);
                $finalControlStates = $finalControlStates | Where-Object { $_.State};
            }
            else
            {
                #merge with the exiting if found
                $persistedControlStates = $this.GetPersistedControlStates("$hash.json");
                $finalControlStates = $this.MergeControlStates($persistedControlStates, $finalControlStates);
                $this.UpdateControlIndexer($id, $finalControlStates, $false);
            }
        }
        else
        {
            #purge would remove the entry from the control indexer and also purge the stale state json.
            $this.PurgeControlState($id);
        }
        [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 {
                $state = $_;
                $loopValue = $this.retryCount;
                while($loopValue -gt 0)
                {
                    $loopValue = $loopValue - 1;
                    try
                    {
                        Set-AzureStorageBlobContent -File $state.FullName -Container $ContainerName -BlobType Block -Context $StorageAccount.Context -Force -ErrorAction Stop
                        $loopValue = 0;
                    }
                    catch
                    {
                        #eat this exception and retry
                    }
                }
            }
        }
        else
        {
            #clean up the container as there is no indexer
            Get-AzureStorageBlob -Container $ContainerName -Context $StorageAccount.Context | Remove-AzureStorageBlob  
        }
    }

    hidden [void] PurgeControlState([string] $id)
    {        
        $AzSDKTemp = [Constants]::AzSdkAppFolderPath + "\Temp\ServerControlState";                
        if(-not (Test-Path "$AzSDKTemp\ControlState"))
        {
            mkdir -Path "$AzSDKTemp\ControlState" -ErrorAction Stop | Out-Null
        }
        else
        {
            Remove-Item -Path "$AzSDKTemp\ControlState\*" -Force 
        }

        $hash = [Helpers]::ComputeHash($id);
        $indexerPath = "$AzSDKTemp\ControlState\$($this.IndexerBlobName)"
        $fileName = "$AzSDKTemp\ControlState\$hash.json"    
        
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = ""
        if($null -ne $this.AzSDKStorageContainer)
        {
            $ContainerName = $this.AzSDKStorageContainer.Name
        }
        
        $this.UpdateControlIndexer($id, $null, $true);
        if($null -ne $this.ControlStateIndexer)
        {                
            [Helpers]::ConvertToJsonCustom($this.ControlStateIndexer) | Out-File $indexerPath -Force
            $controlStateArray = Get-ChildItem -Path "$AzSDKTemp\ControlState"                
            $controlStateArray | ForEach-Object {
                $state = $_
                $loopValue = $this.retryCount;
                while($loopValue -gt 0)
                {
                    $loopValue = $loopValue - 1;
                    try
                    {
                        Set-AzureStorageBlobContent -File $state.FullName -Container $ContainerName -BlobType Block -Context $StorageAccount.Context -Force -ErrorAction Stop
                        $loopValue = 0;
                    }
                    catch
                    {
                        #eat this exception and retry
                    }
                }
            }
        }
        $loopValue = $this.retryCount;
        while($loopValue -gt 0)
        {
            $loopValue = $loopValue - 1;
            try
            {
                Remove-AzureStorageBlob -Blob "$hash.json" -Context $StorageAccount.Context -Container $ContainerName -Force -ErrorAction Stop
                $loopValue = 0;
            }
            catch
            {
                #eat this exception and retry
            }
        }        
    }

    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 = ""
        if($null -ne $this.AzSDKStorageContainer)
        {
            $ContainerName = $this.AzSDKStorageContainer.Name
        }
        [ControlState[]] $ControlStatesJson = @()

        $loopValue = $this.retryCount;
        while($loopValue -gt 0)
        {
            $loopValue = $loopValue - 1;
            try
            {
                $controlStateBlob = Get-AzureStorageBlob -Container $ContainerName -Blob $controlStateBlobName -Context $StorageAccount.Context -ErrorAction Stop
                Get-AzureStorageBlobContent -CloudBlob $controlStateBlob.ICloudBlob -Context $StorageAccount.Context -Destination "$AzSDKTemp\ExistingControlStates" -Force -ErrorAction Stop
                $ControlStatesJson = Get-ChildItem -Path "$AzSDKTemp\ExistingControlStates\$controlStateBlobName" -Force -ErrorAction Stop | Get-Content | ConvertFrom-Json 
                $loopValue = 0;
            }
            catch
            {
                $ControlStatesJson = @()
                #eat this exception and retry
            }
        }
            
        if(($ControlStatesJson | Measure-Object).Count -gt 0)
        {
            $ControlStatesJson | ForEach-Object {
                $ControlState = $_;
                #this can be removed after we check there is no value for attestationstatus coming to azsdktm database
                if($ControlState.AttestationStatus -eq [AttestationStatus]::NotFixed)
                {
                    $ControlState.AttestationStatus = [AttestationStatus]::WillNotFix
                }
            }
        }

        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) -and ($_.ChildResourceName -eq $controlState.ChildResourceName) } | 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, [bool] $ToBeDeleted)
    {
        $this.ControlStateIndexer = $null;
        $retVal = $this.ComputeControlStateIndexer();
        $StorageAccount = $this.AzSDKStorageAccount;                        
        $containerObject = $this.AzSDKStorageContainer
        $ContainerName = ""
        if($null -ne $this.AzSDKStorageContainer)
        {
            $ContainerName = $this.AzSDKStorageContainer.Name
        }
        if($retVal)
        {                
            $tempHash = [Helpers]::ComputeHash($id);            
            $filteredIndexerObject = $this.ControlStateIndexer | Where-Object { (($_.HashId -eq $tempHash -and -not $ToBeDeleted) -or ($ToBeDeleted -and $_.HashId -ne $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
        }

    }

    
}