<#
    .SYNOPSIS
        Connect to an Azure VM via RDP, starting it if it is not already running

    .DESCRIPTION
        VM credentials are cached to Credential Manager and won't be required after the first connect

	.PARAMETER Subscription
        The Azure Subscription that contains the VMS
    .PARAMETER ResourceGroupName
        The Azure Resource Group that contains the VM
    .PARAMETER VMName
        The name of the VM to connect to

    .EXAMPLE
        Connect-AzureVmRDP -Subscription "SUB1" -ResourceGroup "dev" -VMName "MyVM"

    .FUNCTIONALITY
        Azure
#>
function Connect-AzureVmRDP
{
	param(
		[string]$Subscription,
		[string]$ResourceGroupName,
		[string]$VMName
	)

    $vm = Start-AzureRmVM -Subscription $Subscription -ResourceGroupName $ResourceGroupName -VMName $VMName

	Write-Host "Finding IP Address"
	$nic = Get-AzureRmNetworkInterface | ? { $_.VirtualMachine.Id -eq $vm.Id }

	$ipId = $nic.IpConfigurations[0].PublicIpAddress.Id

    $ipAddress = (Get-AzureRmPublicIpAddress | ?{ $_.Id -eq $ipId }).IpAddress

	$publicIp = Get-AzureRmPublicIpAddress | ?{ $_.Id -eq $nic.IpConfigurations[0].PublicIpAddress.Id } | select -First 1

	$vmHost = $publicIp.DnsSettings.Fqdn

	if (-not $vmHost)
	{
		$vmHost = $publicIp
	}

	Write-Host "Waiting for RDP port 3349 to be open on $vmHost"
    if (-not (Wait-TCPPort -Server $vmHost -Port 3389 -TimeoutSeconds 120))
    {
        throw "Timed out waiting for RDP to be available on $vmHost"
    }

	Write-Host "Checking for stored credentials"
    Assert-StoredCredential $vmHost

	Write-Host "Connecting to RDP session"
    & "start" "mstsc" "/v:$vmHost" | Out-Null

	# IPs can change, so if we don't have public dns just remove the credential entry rather than filling up the vault with temporary IPs
	if ($vmHost -eq $publicIp)
	{
		Remove-StoredCredential $vmHost
	}
}

function Start-AzureRmVM
{
	param(
		[string]$Subscription,
		[string]$ResourceGroupName,
		[string]$VMName
	)

    if ((Get-AzureRmContext).Subscription.Name -ne $Subscription)
    {
		Write-Host "Logging into Azure subscription '$Subscription'"

        Login-AzureRmAccount | Out-Null
        Select-AzureRmSubscription -Subscription $Subscription | Out-Null
    }

	Write-Host "Finding VM '$VMName'"

    $vm = Get-AzureRmVM -ResourceGroupName $ResourceGroupName -Name $VMName -Status

    if (-not ($vm | ?{ $_.Statuses.Code -eq "PowerState/running" }))
    {
		Write-Host "Starting VM '$VMName'"

        Start-AzureRmVM -ResourceGroupName $ResourceGroupName -Name $VMName | Out-Null
    }

    return (Get-AzureRmVM -ResourceGroupName $ResourceGroupName -Name $VMName)
}

function Assert-StoredCredential
{
    param($name)

    if (Test-GenericCredential -name $name)
    {
        return
    }

    $credential = Get-Credential -Message $name

    if (-not $credential)
    {
        throw "Cancelled"
    }

    $networkCredential = $credential.GetNetworkCredential()

    & 'cmdkey' "/generic:$name" "/user:$($networkCredential.UserName)" "/password:`"$($networkCredential.Password)`"" | Out-Null
}

function Remove-StoredCredential
{
	param($name)

	& 'cmdkey' "/delete:$name" | Out-Null
}

function Test-GenericCredential {
    Param(
    [string]$name
    )

    $cmd = (& 'cmdkey' "/list:$name" | Select-Object -Skip 3)

    return $cmd -ne "* NONE *"
}

function Test-TCPConnection
{
    param(
        [string]$Server,
        [int]$Port
    )


    $Connection = New-Object System.Net.Sockets.TCPClient

    try {
        $Connection.Connect($Server,$Port)
        $Connection.Close()
        return $true
    }
    catch {
        return $false
    }
}

function Wait-TCPPort
{
    param(
        [string]$Server,
        [int]$Port,
        [int]$TimeoutSeconds
    )

    $TotalWaited = 0
    $SleepSeconds = 2

    while ($TotalWaited -lt $TimeoutSeconds)
    {
        if (Test-TCPConnection -Server $Server -Port $Port) {
            return $true
        }

        Start-Sleep -Seconds $SleepSeconds
        $TotalWaited += $SleepSeconds
    }

    return $false
}