Write-Disk.ps1

<#PSScriptInfo
 
.VERSION 1.0.1
 
.GUID 530d799e-70a6-4c03-905e-d316933c9e60
 
.AUTHOR Sonny Sasaka
 
.COPYRIGHT (c) 2025 Sonny Sasaka. All rights reserved.
 
.TAGS disk image iso usb sd write raw physical
 
.LICENSEURI https://opensource.org/licenses/MIT
 
.RELEASENOTES
Initial release of Write-Disk script for writing disk images to physical drives.
 
#>


<#
.SYNOPSIS
    Writes a disk image file to a physical drive.
 
.DESCRIPTION
    This script writes a raw disk image (e.g., ISO, IMG) to a physical drive such as a USB flash drive or SD card.
    By default, it only allows writing to removable drives (USB, SD, MMC) to prevent accidental data loss.
    Requires Administrator privileges.
 
.PARAMETER ImageFile
    Path to the disk image file to write (e.g., ubuntu.iso, raspbian.img).
 
.PARAMETER DiskNumber
    Target disk number (e.g., 1 for Disk 1). Use -ListDisks to see available disks.
 
.PARAMETER BufferSize
    Buffer size for read/write operations. Default is 1MB.
 
.PARAMETER Force
    Bypass the removable drive safety check. Use with caution.
 
.PARAMETER ListDisks
    List all available physical disks and exit.
 
.EXAMPLE
    .\Write-Disk.ps1 -ListDisks
    Lists all available physical disks.
 
.EXAMPLE
    .\Write-Disk.ps1 -ImageFile "ubuntu.iso" -DiskNumber 1
    Writes ubuntu.iso to Disk 1 (which must be a removable drive).
 
.EXAMPLE
    .\Write-Disk.ps1 -ImageFile "ubuntu.iso" -DiskNumber 0 -Force
    Writes to Disk 0 even if it's not a removable drive (use with extreme caution).
 
.NOTES
    Author: Sonny Sasaka
    Requires: Administrator privileges
    Safety: Only writes to removable drives by default (USB, SD, MMC)
#>

param(
    [Parameter(Mandatory=$false)]
    [string]$ImageFile,

    [Parameter(Mandatory=$false)]
    [int]$DiskNumber = -1,

    [int]$BufferSize = 1MB,

    [switch]$Force,

    [switch]$ListDisks
)

# Resolve ImageFile to absolute path before admin elevation changes working directory
if ($ImageFile -and -not [System.IO.Path]::IsPathRooted($ImageFile)) {
    $ImageFile = Join-Path (Get-Location).Path $ImageFile
}

function Show-AvailableDisks {
    Write-Host "Available physical disks:"
    try {
        Get-Disk | Sort-Object Number
    } catch {
        Write-Host " (Unable to enumerate disks)"
    }
}

# Writing to raw physical disk requires Administrator privilege
$currentPrincipal = [Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()
if (-not $currentPrincipal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
    Write-Error "This script must be run as Administrator"
    exit 1
}

# Handle -ListDisks flag
if ($ListDisks) {
    Show-AvailableDisks
    exit 0
}

# Validate required parameters
if (-not $ImageFile) {
    Write-Error "ImageFile parameter is required"
    exit 1
}
if ($DiskNumber -lt 0) {
    Write-Error "DiskNumber parameter is required"
    exit 1
}

# Validate image file exists
if (-not (Test-Path $ImageFile)) {
    Write-Error "Image file not found: $ImageFile"
    exit 1
}

# Get image file info
$imageInfo = Get-Item $ImageFile
$totalBytes = $imageInfo.Length
$totalMB = [math]::Round($totalBytes / 1MB, 2)

# Validate we can actually open the image file before proceeding
try {
    $testStream = [System.IO.File]::OpenRead($ImageFile)
    $testStream.Close()
} catch {
    Write-Error "Cannot open image file: $_"
    exit 1
}

# Get disk information
try {
    $diskInfo = Get-Disk -Number $DiskNumber -ErrorAction SilentlyContinue
} catch {
    # Ignore errors, diskInfo will remain null
}

# Validate disk exists
if (-not $diskInfo) {
    Write-Error "Disk $DiskNumber does not exist."
    Write-Host ""
    Show-AvailableDisks
    exit 1
}

# Check if disk is removable
$isRemovable = $diskInfo.BusType -eq 'USB' -or $diskInfo.BusType -eq 'SD' -or $diskInfo.BusType -eq 'MMC'

if (-not $isRemovable) {
    if (-not $Force) {
        Write-Error "Safety check failed: Disk $DiskNumber is not a removable drive (BusType: $($diskInfo.BusType))"
        Write-Error "This script only works with removable drives (USB, SD, MMC) to prevent accidental data loss."
        Write-Error "Use -Force parameter to override this check (if you know what you are doing)."
        exit 1
    } else {
        Write-Warning "FORCE MODE: Bypassing removable drive check!"
        Write-Warning "Disk $DiskNumber is not a removable drive (BusType: $($diskInfo.BusType))"
        Write-Warning "Proceeding with caution..."
    }
}

$diskSizeGB = [math]::Round($diskInfo.Size / 1GB, 2)
$diskModel = if ($diskInfo.FriendlyName) { $diskInfo.FriendlyName } else { "Unknown" }
$busType = $diskInfo.BusType

Write-Host "Source: $ImageFile ($totalMB MB)"
Write-Host "Target: Disk $DiskNumber"
Write-Host " Model: $diskModel"
Write-Host " Size: $diskSizeGB GB"
Write-Host " Bus Type: $busType"
Write-Host ""
Write-Warning "ALL DATA ON DISK $DiskNumber WILL BE DESTROYED!"
$confirm = Read-Host "Type 'YES' to continue"

if ($confirm -ne 'YES') {
    Write-Host "Cancelled."
    exit 0
}

# Clean the disk
Write-Host ""
Write-Host "Cleaning disk $DiskNumber..."

try {
    Clear-Disk -Number $DiskNumber -RemoveData -RemoveOEM -Confirm:$false -ErrorAction Stop
} catch {
    Write-Warning "Clear-Disk encountered an issue: $_"
    Write-Warning "Continuing anyway..."
}

Start-Sleep -Seconds 1

Write-Host ""
Write-Host "Writing image..."

try {
    # Open source file for reading
    $sourceStream = [System.IO.File]::OpenRead($ImageFile)

    # Open physical drive for writing
    $physicalDrivePath = "\\.\PhysicalDrive$DiskNumber"
    $driveStream = [System.IO.File]::Open(
        $physicalDrivePath,
        [System.IO.FileMode]::Open,
        [System.IO.FileAccess]::Write,
        [System.IO.FileShare]::None
    )

    $buffer = New-Object byte[] $BufferSize
    $bytesWritten = 0
    $lastPercent = -1

    while ($true) {
        $bytesRead = $sourceStream.Read($buffer, 0, $buffer.Length)
        if ($bytesRead -eq 0) { break }

        $driveStream.Write($buffer, 0, $bytesRead)
        $bytesWritten += $bytesRead

        # Simple progress indicator
        $percent = [math]::Floor(($bytesWritten / $totalBytes) * 100)
        $writtenMB = [math]::Round($bytesWritten / 1MB, 2)

        if ($percent -ne $lastPercent) {
            Write-Host "`r$percent% - $writtenMB / $totalMB MB" -NoNewline
            $lastPercent = $percent
        }
    }

    Write-Host ""
    Write-Host "Flushing buffers..."
    $driveStream.Flush()

    Write-Host "Done! $totalMB MB written successfully."

} catch {
    Write-Error "Error: $_"
    exit 1
} finally {
    if ($sourceStream) { $sourceStream.Close() }
    if ($driveStream) { $driveStream.Close() }
}