customData.ps1

$transcriptsDirectory = "C:\transcripts"
$currentUser = [Security.Principal.WindowsIdentity]::GetCurrent()
$currentUserName = $currentUser.Name
$transcriptUser = $currentUserName.replace('\', '_')

if (!(Test-Path $transcriptsDirectory)) { New-Item -Path $transcriptsDirectory -ItemType Directory }
if ($currentUserName.Contains("SYSTEM")) { icacls "C:\AzureData" /grant Everyone:F /T }

$date = Get-Date
$transcriptFileName = "customData" + '_' + $transcriptUser + '_' + $date.ToString("yyyy-MM-dd_HH-mm-ss") + ".log"
# Start the transcript
Start-Transcript -Path $transcriptsDirectory\$transcriptFileName

### Define functions
function Get-InstanceMetadata {
    $metadataUrl = 'http://169.254.169.254/metadata/'
    $url = $metadataUrl + '/instance?api-version=2021-01-01&format=json'
    $headers = @{Metadata = 'true' }
    $instanceMetadata = Invoke-RestMethod -Uri $url -Headers $headers -Method GET -UseBasicParsing
    return $instanceMetadata
}

function Get-Tags {
    $instanceMetadata = Get-InstanceMetadata
    $instanceMetadata = $instanceMetadata.compute.tags
    $pairs = $instanceMetadata -split '; ?'
    $tags = New-Object PSObject
    foreach ($pair in $pairs) {
        $key, $value = $pair -split ':', 2
        $tags | Add-Member -NotePropertyName $key -NotePropertyValue $value
    }
    return $tags
}

function Get-AccessToken {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory)]
        [string] $Audience
    )
    $metadataUrl = 'http://169.254.169.254/metadata/'
    $url = $metadataUrl + '/identity/oauth2/token?api-version=2018-02-01&resource=' + $Audience
    $headers = @{Metadata = 'true' }
    $accessToken = Invoke-RestMethod -Uri $url -Headers $headers -Method GET -UseBasicParsing | Select-Object -ExpandProperty access_token
    return $accessToken
}

function Get-StorageAccountKey {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory)]
        [string] $StorageAccountResourceId
    )
    $accessToken = Get-AccessToken -Audience 'https://management.azure.com/'
    $keys = Invoke-RestMethod -Uri "https://management.azure.com/$StorageAccountResourceId/listKeys?api-version=2022-09-01&$expand=kerb" -Headers @{
        'Accept'        = 'application/json'
        'Authorization' = 'Bearer ' + $accessToken
    } -Method POST -Body '' | Select-Object -ExpandProperty keys
    $key = $keys | Where-Object { $_.keyName -eq 'key1' } | Select-Object -ExpandProperty value
    return $key
}

function Connect-FileShare {
    param (
        [Parameter(Mandatory)]
        [string] $StorageAccountName,
        [Parameter(Mandatory)]
        [string] $FileShareName,
        [Parameter(Mandatory)]
        [string] $AccountKey,
        [Parameter(Mandatory)]
        [string] $DriveLetter
    )
    $filesShareUrl = $StorageAccountName + '.file.core.windows.net'
    $connectTestResult = Test-NetConnection -ComputerName $filesShareUrl -Port 445
    if ($connectTestResult.TcpTestSucceeded) {
        cmd.exe /C "cmdkey /add:`"$filesShareUrl`" /user:`"localhost\$StorageAccountName`" /pass:`"$AccountKey`""
        $driveExists = Get-PSDrive -Name $DriveLetter -ErrorAction SilentlyContinue
        if ($driveExists) { Remove-PSDrive -Name $DriveLetter -Force }
        $fileShareUncPath = '\\' + $filesShareUrl + '\' + $FileShareName
        New-PSDrive -Name $DriveLetter -PSProvider FileSystem -Root "$fileShareUncPath" -Persist -Scope Global -ErrorAction SilentlyContinue
    }
    else {
        Write-Error -Message "Unable to reach the Azure storage account via port 445."
    }
}
### End of functions

# CustomScriptExtensionActions
if ($currentUserName.Contains("SYSTEM")) {

    $RegistryPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Run"
    $Name = "CustomData"
    $Value = "powershell.exe -File `"C:\AzureData\CustomData.ps1`""

    if (!(Test-Path $RegistryPath)) { New-Item -Path $RegistryPath -Force | Out-Null }
    $PropertyExists = Get-ItemProperty -Path $RegistryPath -Name $Name -ErrorAction SilentlyContinue | Out-Null
    if (-not $PropertyExists) {
        New-ItemProperty -Path $RegistryPath -Name $Name -Value $Value -PropertyType String -Force | Out-Null
    }

}
else {
# Startup script actions
    $tags = Get-Tags
    $storageAccountResourceId = $tags.storageAccountResourceId
    $storageAccountName = $storageAccountResourceId.split("/")[-1]
    $fileshareName = $tags.fileShareName
    $accountKey = Get-StorageAccountKey -StorageAccountResourceId $storageAccountResourceId

    Connect-FileShare -StorageAccountName $storageAccountName -FileShareName $fileshareName -AccountKey $accountKey -DriveLetter Z
}


###
Stop-Transcript