
param(
    [Parameter(Mandatory = $true)]
    [string]$VMName,
    [Parameter(Mandatory = $false)]
    [int]$diskNumber = 1, <# 0-based indexing #>
    [Parameter(Mandatory = $true)]
    [string]$AdminPassword
)

$password = ConvertTo-SecureString $AdminPassword -AsPlainText -Force
$cred = New-Object System.Management.Automation.PSCredential ("Administrator", $password)

Write-Host "Erasing connected disk's readonly attribute"
$diskpartScriptFindVolume = @"
select disk $diskNumber
detail disk
"@
$diskpartScriptDisableReadonly = @"
select disk $diskNumber
attributes disk clear readonly
"@


function RemoteExecuteDiskpartScript([string]$diskpartScript)
{
    $output = Invoke-Command -VMName $VMName -Credential $cred -ScriptBlock {
        param($diskpartScript)
    
        $diskpartScriptPath = "C:\AOFG\diskpartScript.txt"
        $diskpartScriptOutput = "C:\AOFG\diskpartOutput.txt"
        Set-Content -Path $diskpartScriptPath -Value $diskpartScript
        Start-Process -FilePath "diskpart.exe" -ArgumentList "/s $diskpartScriptPath" -Wait -RedirectStandardOutput $diskpartScriptOutput
        
        $output = Get-Content -Path $diskpartScriptOutput

        # Clean up the diskpart script file
        Remove-Item -Path $diskpartScriptPath
        Remove-Item -Path $diskpartScriptOutput

        return $output
    } -ArgumentList $diskpartScript
    return $output
}

RemoteExecuteDiskpartScript -diskpartScript $diskpartScriptDisableReadonly | Out-Null
$rawDiskInfo = RemoteExecuteDiskpartScript -diskpartScript $diskpartScriptFindVolume
$pattern = "Volume\s+(\d+)"
$volumeNumber = ($rawDiskInfo | Select-String -Pattern $pattern)
$volumeNumber = $volumeNumber.Matches.Groups[1].Value

$diskpartScriptAssignAsFolder = @"
select disk $diskNumber
select vol $volumeNumber
remove all noerr
assign mount="D:"
"@

RemoteExecuteDiskpartScript -diskpartScript $diskpartScriptAssignAsFolder
