Framework/Core/InClusterCA/ContinuousAssuranceForHDInsight.ps1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
class HDInsightClusterCA : CommandBase {
    [PSObject] $ResourceContext;

    HDInsightClusterCA([PSObject] $ResourceContext, [InvocationInfo] $invocationContext): 
    Base([Constants]::BlankSubscriptionId, $invocationContext) { 
        $this.ResourceContext = $ResourceContext
    }

    static [PSObject] GetParameters($SubscriptionId, $ClusterName, $ResourceGroupName) {
        if([string]::IsNullOrEmpty($SubscriptionId) -or 
            [string]::IsNullOrEmpty($ClusterName) -or
            [string]::IsNullOrEmpty($ResourceGroupName)) {

            Write-Host "Input the following parameters"
            if ([string]::IsNullOrEmpty($SubscriptionId)) {
                $SubscriptionId = [Helpers]::ReadInput("Subscription ID")
            }
            if ([string]::IsNullOrEmpty($ClusterName)) {
                $ClusterName = [Helpers]::ReadInput("HDInsight Cluster Name")
            }
            if ([string]::IsNullOrEmpty($ResourceGroupName)) {
                $ResourceGroupName = [Helpers]::ReadInput("HDInsight Resource Group Name")
            }
        }
        Set-AzContext -SubscriptionId $SubscriptionId *> $null
        $Cluster = Get-AzHDInsightCluster -ClusterName $ClusterName -ErrorAction Ignore
        if ($null -eq $Cluster) { 
            Write-Host "HDInsight cluster [$ClusterName] wasn't found. Please retry" -ForegroundColor Red
            throw $_;
        }
        $ResourceGroup = $cluster.ResourceGroup
        # Extracting Storage Account
        $StorageAccount = $cluster.DefaultStorageAccount.Split(".")[0]
        $StorageAccountContext = Get-AzStorageAccount -Name $storageAccount `
                                    -ResourceGroupName $ResourceGroup `
                                    -ErrorAction Ignore
        while ($null -eq $StorageAccountContext) {
            Write-Host "Storage [$StorageAccount] not found in resource group [$resourceGroup]."
            $NewRGName = Read-Host -Prompt "Enter the resource group name where [$storageAccount] is present."
            $StorageAccountContext = Get-AzStorageAccount -Name $StorageAccount -ResourceGroupName $NewRGName
        }
        return @{
            "SubscriptionId" =          $SubscriptionId;
            "StorageAccountContext" =   $StorageAccountContext.Context;
            "Container" =               $Cluster.DefaultStorageContainer;
            "ResourceGroup" =           $Cluster.ResourceGroup;
            "ClusterName" =             $Cluster.Name;
        }
    }


    [void] UploadAzSKNotebookToCluster() {
        $NotebookUrl = [Constants]::HDInsightCANotebookUrl
        $FilePath = $env:TEMP + "\AzSK_CA_Scan_Notebook.ipynb"
        Invoke-RestMethod  -Method Get -Uri $NotebookUrl -OutFile $filePath
        (Get-Content $FilePath) -replace '\#AI_KEY\#', $this.ResourceContext.InstrumentationKey | Set-Content $FilePath -Force
        (Get-Content $FilePath) -replace '\#SID\#', $this.ResourceContext.SubscriptionId | Set-Content $FilePath -Force
        (Get-Content $FilePath) -replace '\#RG_NAME\#', $this.ResourceContext.ResourceGroup | Set-Content $FilePath -Force
        (Get-Content $FilePath) -replace '\#RES_NAME\#', $this.ResourceContext.ClusterName | Set-Content $FilePath -Force
        Set-AzStorageBlobContent -File $FilePath -Blob "HdiNotebooks/PySpark/AzSK_CA_Scan_Notebook.ipynb" `
                                -Container $this.ResourceContext.Container `
                                -Context $this.ResourceContext.StorageAccountContext | Out-Null
        Remove-Item $FilePath -ErrorAction Ignore
    }

    [void] RemoveAzSKNotebookFromCluster() {
        Remove-AzStorageBlob -Blob "HdiNotebooks/PySpark/AzSK_CA_Scan_Notebook.ipynb" `
                             -Container $this.ResourceContext.Container `
                             -Context $this.ResourceContext.StorageAccountContext
    }

    [void] InstallAzSKPy() {
        $ScriptActionUri = [Constants]::AzSKPyInstallUrl
        # Install on both head, and worker nodes
        $NodeTypes = "headnode", "workernode"
        $ScriptActionName = "AzSKInstaller-" + (Get-Date -f MM-dd-yyyy-HH-mm-ss)
        Submit-AzHDInsightScriptAction -ClusterName $this.ResourceContext.clusterName `
                                        -Name $ScriptActionName `
                                        -Uri $ScriptActionUri `
                                        -NodeTypes $NodeTypes `
                                        -PersistOnSuccess > $null
     
    }

    [void] UninstallAzSKPy() {
        $uninstallScript = [Constants]::AzSKPyUninstallUrl
        # Uninstall on both head, and worker nodes
        $nodeTypes = "headnode", "workernode"
        $uninstallActionName = "AzSKUnInstaller-" + (Get-Date -f MM-dd-yyyy-HH-mm-ss)
        $this.PublishCustomMessage("Uninstalling AzSKPy on the cluster")
        Submit-AzHDInsightScriptAction -ClusterName $this.ResourceContext.ClusterName `
                                            -Name $uninstallActionName `
                                            -Uri $uninstallScript `
                                            -NodeTypes $nodeTypes `
                                            -PersistOnSuccess > $null
    }

    [void] InstallCA() {
        $this.PublishCustomMessage("Uploading Notebook to the cluster")
        $this.UploadAzSKNotebookToCluster()
        $this.PublishCustomMessage("Installing AzSK python library in the cluster")
        $this.InstallAzSKPy()
    }

    [void] RemoveCA() {
       # Submit uninstall script
       $this.UninstallAzSKPy()
       # Remove notebook
       $this.RemoveAzSKNotebookFromCluster()
    }

    [void] UpdateCA() {
        $this.PublishCustomMessage("Updating scan Notebook on the cluster")
        $this.UploadAzSKNotebookToCluster()
        $this.PublishCustomMessage("Updating AzSK python library in the cluster")
        # uninstall first
        $this.UninstallAzSKPy()
        # install new version
        $this.InstallAzSKPy()       
    }

    [void] GetCA() {
        $metapathtemp = $env:TEMP + "\azskmetatemp.json"
        $metapath = $env:TEMP + "\azskmeta.json"
        $notebook = Get-AzStorageBlob -Blob "HdiNotebooks/PySpark/AzSK_CA_Scan_Notebook.ipynb" `
                                      -Container $this.ResourceContext.Container `
                                      -Context $this.ResourceContext.StorageAccountContext -ErrorAction Ignore
        if ($notebook -eq $null) {
            $this.PublishCustomMessage("CA Health not OK. Either the installation is broken or not present. Please re-install using Install-AzSKContinuousAssuranceForCluster", [MessageType]::Error)
            return;
        }
        $list = New-Object System.Collections.Generic.List[System.Object]
        $filesList = Get-AzStorageBlob -Blob "" -Container $this.ResourceContext.Container -Context $this.ResourceContext.StorageAccountContext
        foreach ($x in $filesList.Name) {
            if ($x.Contains("AzSK_Meta") -and $x.Contains("part") -and $x.Contains("json")) {
                $content = Get-AzStorageBlob -Blob $x -Container $this.ResourceContext.Container -Context $this.ResourceContext.StorageAccountContext
                if ($content.Length -gt 0) {
                    $list.Add($content)
                }
            }
        }

        if ($list.Count -eq 0) {
            $this.PublishCustomMessage("Required information not found. Please check if AzSK CA is installed on the cluster", [MessageType]::Error)
        } else {
            $sortedList = $list | Sort LastModified -Descending
            $res = Get-AzStorageBlobContent -Blob $sortedList[0].Name -Container $this.ResourceContext.Container -Context $this.ResourceContext.StorageAccountContext -Destination $metapath -Force
            $json = (Get-Content -Path $metapath)
            $json = $json | ConvertFrom-Json
            $this.PublishCustomMessage("CA Health OK. Following is the summary", [MessageType]::Update)
            $this.PublishCustomMessage([Helpers]::ConvertObjectToString($json, $true))
        }        
    }
}