Startupscript for PowerShell Cloud Shell 2017/10/09

PSCloudShellStartup.ps1

# Startupscript for PowerShell Cloud Shell

# User profile script location
$script:CurrentHostProfilePath = (Join-Path $env:USERPROFILE 'CloudDrive\Microsoft.PowerShell_profile.ps1')
$script:AllHostsProfilePath    = (Join-Path $env:USERPROFILE 'CloudDrive\profile.ps1')

$script:AzureRMProfileModule = 'AzureRM.Profile'
if ($PSVersionTable.PSEdition -eq 'Core')
{
    $script:AzureRMProfileModule = 'AzureRM.Profile.NetCore'
}

#region Utility Functions

function Login-PSCloudShellAzureService
{   
    # Clean up completed PSCloudShell jobs
    Get-Job -Name PSCloudShell* | Where-Object {$_.State -eq 'Completed'} | Remove-Job -Force -ErrorAction SilentlyContinue

    $acContent = Microsoft.PowerShell.Management\Get-Content $psAccessTokensFilePath | out-string
    $acContent = $acContent.Replace('"','').Trim()

    # File contains Azure auth token in this format - "AccessToken;KeyVaultToken;GraphToken"
    # Extract the AccessToken and initialize AZURE_CONSOLE_TOKENS - used for Azure Cli & AzureRM logins
    # Extract the GraphToken and initialize PSGraphToken -used for Azure AD login
    $tokenArray = $acContent.Split(';')
    $env:AZURE_CONSOLE_TOKENS = $tokenArray[0]
    $env:PSGraphToken = $tokenArray[1]

    <#
    # Authenticate to Azure Cli
    $jobName = 'PSCloudShellAzureAuth-' + [guid]::NewGuid()

    # Need to set AZURE_CONFIG_DIR and AZURE_ACCESS_TOKEN_FILE before calling az login so Azure CLI config and logs will be saved in cloud
    if (Test-Path -Path (Join-Path -Path $env:USERPROFILE -ChildPath '\CloudDrive'))
    {
        $azureAccessTokenFile = Join-Path -Path $env:USERPROFILE -ChildPath  '.azure\accessTokens.json'
        [Environment]::SetEnvironmentVariable('AZURE_ACCESS_TOKEN_FILE', $azureAccessTokenFile , 'Process')

        $azureConfigDir = Join-Path -Path $env:USERPROFILE -ChildPath '\CloudDrive\.pscloudshell\.azure'
        [Environment]::SetEnvironmentVariable('AZURE_CONFIG_DIR', $azureConfigDir , 'Process')
    }

    $null = Start-Job -Name $jobName -ScriptBlock `
    {
        # Login to Azure using Access Token - to be used by az cli
        $null = & "${env:ProgramFiles(x86)}\Microsoft SDKs\Azure\CLI\wbin\azure" login

        # Enable azure cli telemetry
        $null = & "${env:ProgramFiles(x86)}\Microsoft SDKs\Azure\CLI\wbin\azure" telemetry --enable
    }
    #>

    # Enable AzureRM Data collection
    # Else User is prompted when Add-AzureRMAccount is invoked
    Set-PSCloudShellTelemetry

    # Authenticate to AzureRM
    Import-Module $script:AzureRMProfileModule
    $addAzureRMAccountParameters = @{'AccessToken' = $env:AZURE_CONSOLE_TOKENS; 'AccountID' = $env:ACC_OID}    

    if ($env:ACC_CLOUD -eq 'dogfood')
    {
        $addAzureRMAccountParameters.Add('EnvironmentName', $env:ACC_CLOUD)
    }

    # For subsequent logins, maintain user's current subscription
    $currSubscription = (Get-AzureRmContext).Subscription.Name
    if($null -ne $currSubscription) {
        $addAzureRMAccountParameters.Add('SubscriptionName', $currSubscription)
    }

    $azureAccount = & $script:AzureRMProfileModule\Add-AzureRMAccount @addAzureRMAccountParameters -ErrorAction SilentlyContinue -ErrorVariable azureRMError

    # Log any errors from AzureRM authentication
    if ($azureRMError)
    {
        $errorFolderPath = $env:USERPROFILE
        if (Test-Path (Join-Path $env:USERPROFILE '.azure'))
        {
            $errorFolderPath = (Join-Path $env:USERPROFILE '.azure')
        }
        $azureRMError > (Join-Path $errorFolderPath 'azureRMError.err')
    }

    # AzureAD is currently not supported on PowerShell Core
    if ($PSVersionTable.PSEdition -eq 'Desktop') {
        # Authenticate to AzureAD

        # Import the Azure AD Auth dll
        Import-Module PSCloudShellADAuth -Force

        # Change the authentication factory
        [Microsoft.Azure.PSCloudConsole.ADAuth.ADAuthFactory]::Update()
    }
    return $azureAccount
}

function Set-PSCloudShellTelemetry
{
    # Default value in case PSCloudShellUtility is not loaded
    Import-Module $script:AzureRMProfileModule
    $productVersion = '0.1.0'
    $productName = 'PSCloudShell'
    $module = Get-Module -Name PSCloudShellUtility -ListAvailable -ErrorAction SilentlyContinue
    if($module)
    {
        $productVersion = $module.Version.ToString()
    }

    [Microsoft.Azure.Common.Authentication.AzureSession]::ClientFactory.AddUserAgent($productName, $productVersion)
    & $script:AzureRMProfileModule\Enable-AzureRMDataCollection -WarningAction SilentlyContinue
}

function Invoke-PSCloudShellUserProfile
{
    # First run all hosts profile
    if(Microsoft.PowerShell.Management\Test-Path $script:AllHostsProfilePath)
    {
        try
        {
            Microsoft.PowerShell.Utility\Write-Verbose -Message 'Loading AllHosts profile ...' -Verbose

            # As the startupscript.ps1 gets executed with "." for global scope, we use "." here to use the startupscript's scope, i.e., global.
            . $script:AllHostsProfilePath
        }
        catch
        {
            # Log a warning and continue if encountering any terminating errors from the running user profile
            Write-Warning -Message "$_"
        }
    }

    # Second run current host profile
    if(Microsoft.PowerShell.Management\Test-Path $script:CurrentHostProfilePath)
    {
        try
        {
            Microsoft.PowerShell.Utility\Write-Verbose -Message 'Loading CurrentHost profile ...' -Verbose

            # As the startupscript.ps1 gets executed with "." for global scope, we use "." here to use the startupscript's scope, i.e., global.
            . $script:CurrentHostProfilePath
        }
        catch
        {
            # Log a warning and continue if encountering any terminating errors from the running user profile
            Write-Warning -Message "$_"
        }
    }
}

#endregion

#region Initialization

# Set the user profile path to clouddrive
Set-Variable PROFILE $script:CurrentHostProfilePath -Scope Global
$PROFILE = $PROFILE | Add-Member -MemberType NoteProperty -Name CurrentUserAllHosts -Value $script:AllHostsProfilePath -PassThru
$PROFILE = $PROFILE | Add-Member -MemberType NoteProperty -Name CurrentUserCurrentHost -Value $script:CurrentHostProfilePath -PassThru

# Create a function wrapper for emacs - to start emacs in no window mode in PSCloudShell environment
if (Test-Path (Join-Path $env:ProgramFiles 'emacs\bin\emacs.exe'))
{    
    $null = function global:emacs { & "$env:ProgramFiles\emacs\bin\emacs.exe" -nw $args}
}

# Access token file location
$psAccessTokensFilePath = (Join-Path $env:USERPROFILE '.azure\psAccessTokens.txt')

if (-not (Test-Path $psAccessTokensFilePath))
{
    Microsoft.PowerShell.Utility\Write-Warning -Message "AccessTokens file does not exist - $psAccessTokensFilePath. Skipping authenticating to Azure Services"
    Set-PSCloudShellTelemetry
    . Invoke-PSCloudShellUserProfile
    return
}

# Dogfood initialization script
if ($env:ACC_CLOUD -eq 'dogfood')
{
    Microsoft.PowerShell.Utility\Write-Warning -Message "You are running in a dogfood environment. Please supply a URI for Azure dogfood environment initialization script."
    $dfEnvInitScriptURI = Microsoft.PowerShell.Utility\Read-Host -Prompt "Supply the URI"
    $dfEnvInitScript = Microsoft.PowerShell.Utility\Invoke-WebRequest -Uri $dfEnvInitScriptURI -UseBasicParsing | ForEach-Object Content
    $null = [ScriptBlock]::Create($dfEnvInitScript).Invoke()
}

# Authenticate to Azure services
Microsoft.PowerShell.Utility\Write-Verbose -Verbose -Message 'Authenticating to Azure ...'
if (-not (Login-PSCloudShellAzureService))
{
    Microsoft.PowerShell.Utility\Write-Warning -Message 'Authenticating to Azure failed.'
    . Invoke-PSCloudShellUserProfile
    return
}

#endregion

#region Register to reauthenticate to Azure services upon receiving an updated token from the agent

$fswObj = New-Object System.IO.FileSystemWatcher
$fswObj.Path = Split-Path $psAccessTokensFilePath
$fswObj.Filter = Split-Path $psAccessTokensFilePath -Leaf
$eventIdentifier = 'PSCloudShell-' + [guid]::NewGuid() 

$null = Microsoft.PowerShell.Utility\Register-ObjectEvent -InputObject $fswObj -EventName 'Changed' -SourceIdentifier $eventIdentifier -Action `
{
    Login-PSCloudShellAzureService
}

#endregion

#region Load User profile if exists

. Invoke-PSCloudShellUserProfile

#endregion

#region Set PSDefaultParameterValues for cmdlets
$PSDefaultParameterValues = @{'Install-Module:Scope' = 'CurrentUser'}
#endregion

#region Load the Azure Provider initialize Azure drive

# TODO: Remove this check once AzurePSDrive supports PowerShell Core
if ($PSVersionTable.PSEdition -eq 'Desktop')
{
    Import-Module AzurePSDrive
    Import-Module PSCloudShellUtility
    Microsoft.PowerShell.Utility\Write-Verbose -Verbose -Message 'Building your Azure drive ...'

    $null = Microsoft.PowerShell.Management\New-PSDrive -Name Azure -PSProvider SHiPS -Root "AzurePSDrive#Azure" -Scope Global
    if(-not $?)
    {
        Microsoft.PowerShell.Utility\Write-Warning -Message 'Something went wrong while creating Azure drive. You can still use this shell to run Azure PowerShell commands.'
    }
    else
    {
        Set-Location Azure:
    }
}

#endregion