Invoke-AzVMScript

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

function Invoke-AzVMScript {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [ResourceGroupCompleter()]
        [string]$ResourceGroupName,

        [Parameter(Mandatory)]
        [ResourceNameCompleter("Microsoft.Compute/virtualMachines", "ResourceGroupName")]
        [string]$VMName,

        [Parameter(Mandatory, ParameterSetName = 'ScriptBlock')]
        [ScriptBlock]$ScriptBlock,

        [Parameter(Mandatory, ParameterSetName = 'File')]
        [string]$ScriptPath,
        [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
            }
        }
    } else {
        $Result
    }

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