Invoke-AzVmScript.ps1

using namespace Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters

function Invoke-AzVmScript {
    <#
        .Synopsis
            Runs a script block or file on a given Azure VM

        .Notes
            This takes a surprisingly long time to run, and is a bit flaky. Any output past 4096 characters is not returned.

        .Example
            Invoke-AzVmScript -ResourceGroupName AUSC1-FRG -VMName AUSC1-FRGMGT1 -ScriptBlock {Get-DnsServerDiagnostics}

        .Example
            Invoke-AzVmScript -ResourceGroupName AUSC1-FRG -VMName AUSC1-FRGMGT1 -ScriptPath .\Tests\Integration\TestToRun.ps1 -Parameter @{IPAddress = '50.135.9.12'}
    #>
    [CmdletBinding()]
    param(
        # Resource Group to find the VM in
        [Parameter(Mandatory)]
        [ResourceGroupCompleter()]
        [string]$ResourceGroupName,

        # VM to run the script on
        [Parameter(Mandatory)]
        [ResourceNameCompleter("Microsoft.Compute/virtualMachines", "ResourceGroupName")]
        [string]$VMName,

        # Script block to run. Is temporarily saved to a local file, and cleaned up afterwards.
        [Parameter(Mandatory, ParameterSetName = 'ScriptBlock')]
        [ScriptBlock]$ScriptBlock,

        # Path to a script to run.
        [Parameter(Mandatory, ParameterSetName = 'File')]
        [string]$ScriptPath,

        # Parameters to pass to the script
        [HashTable]$Parameter
    )
    if ($ScriptBlock) {
        $ScriptPath = (New-Item -Path $env:Temp -Name "$(New-Guid).ps1" -Value $ScriptBlock.ToString()).FullName
    }

    $Parameters = @{
        ResourceGroupName = $ResourceGroupName
        VMName = $VMName
        CommandId = "RunPowerShellScript"
        ScriptPath = $ScriptPath
    }
    if ($Parameter) {
        $Parameters.Parameter = $Parameter
    }

    $Result = Invoke-AzVmRunCommand @Parameters
    if ($Result.Status -eq 'Succeeded') {
        switch ($Result.Value) {
            {$_.Code -eq 'ComponentStatus/StdOut/succeeded'} {
                $_.Message
                if ($_.Message.Length -eq 4096) {
                    Write-Warning "Message is exactly 4096 characters long, and was likely truncated"
                }
            }
            {$_.Code -eq 'ComponentStatus/StdErr/succeeded' -and $_.Message} {
                Write-Error $_.Message -ErrorAction Continue
            }
        }
    } else {
        $Result
    }

    if ($ScriptBlock) {
        Remove-Item $ScriptPath
    }
}