From d07adbb14c27975d931c0a6cf457f4d1c1bb6649 Mon Sep 17 00:00:00 2001
From: Ben Broderick Phillips <bebroder@microsoft.com>
Date: Mon, 3 Jun 2024 15:50:17 -0400
Subject: [PATCH] Restrict live test storage account access to client IP

---
 eng/common/TestResources/New-TestResources.ps1 | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/eng/common/TestResources/New-TestResources.ps1 b/eng/common/TestResources/New-TestResources.ps1
index 64e075c57049..a78d438b4969 100644
--- a/eng/common/TestResources/New-TestResources.ps1
+++ b/eng/common/TestResources/New-TestResources.ps1
@@ -827,6 +827,19 @@ try {
                                                                 -templateFile $templateFile `
                                                                 -environmentVariables $EnvironmentVariables
 
+        $storageAccounts = Retry { Get-AzResource -ResourceGroupName $ResourceGroupName -ResourceType "Microsoft.Storage/storageAccounts" }
+        if ($storageAccounts) {
+            $clientIp = Retry { Invoke-RestMethod -Uri 'https://myip.dnsomatic.com/' }  # OpenDNS owned ip site. Some throttling will occur if called quickly.
+        }
+        foreach ($account in $storageAccounts) {
+            $rules = Get-AzStorageAccountNetworkRuleSet -ResourceGroupName $ResourceGroupName -AccountName $account.Name
+            if ($rules -and $rules.DefaultAction -eq "Allow") {
+                Write-Host "Restricting network rules in storage account '$($account.Name)' to deny access except from the current client's IP"
+                Retry { Update-AzStorageAccountNetworkRuleSet -ResourceGroupName $ResourceGroupName -Name $account.Name -DefaultAction Deny }
+                Retry { Add-AzStorageAccountNetworkRule -ResourceGroupName $ResourceGroupName -Name $account.Name -IPAddressOrRange $clientIp }
+            }
+        }
+
         $postDeploymentScript = $templateFile.originalFilePath | Split-Path | Join-Path -ChildPath "$ResourceType-resources-post.ps1"
         if (Test-Path $postDeploymentScript) {
             Log "Invoking post-deployment script '$postDeploymentScript'"