param(
	[parameter( ValueFromRemainingArguments = $true )]
	[string[]]$Args # Leave all argument validation to the script, not to PowerShell
)

if ( $args.Length -gt 0 ) {
	Clear-Host
	Write-Host
	Write-Host "GetHDDStatus.ps1,  Version 2.03"
	Write-Host "Get the SMART status for all local harddisks"
	Write-Host
	Write-Host "Usage:    " -NoNewline
	if ( $HOME[0] -eq '/' ) {
		# Linux
		Write-Host "pwsh  GetHDDStatus.ps1" -ForegroundColor White
	} else {
		# Windows
		Write-Host "powershell  . GetHDDStatus.ps1" -ForegroundColor White
	}
	Write-Host
	Write-Host "Notes:    Disks' status is shown in the console and as desktop notification."
	Write-Host "          This script requires elevated privileges (a.k.a. `'root access`')."
	Write-Host "          In Linux, however, it should " -NoNewline
	Write-Host "NOT" -ForegroundColor Red -NoNewline
	Write-Host " be started as root if you want"
	Write-Host "          a desktop notification; the script will restart itself as root to"
	Write-Host "          check the disks, and then send the desktop notification with the"
	Write-Host "          initial non-root user account."
	Write-Host "          In Windows the output includes disk indexes, in Linux it does not."
	Write-Host "          In Linux this script requires smartmontools:"
	Write-Host "          https://www.smartmontools.org/" -ForegroundColor DarkGray
	Write-Host
	Write-Host "Credits:  Windows disk check based on code by Geoff @ UVM:"
	Write-Host "          www.uvm.edu/~gcd/2013/01/which-disk-is-that-volume-on" -ForegroundColor DarkGray
	Write-Host "          System Tray ToolTip Balloon code by Don Jones:"
	Write-Host "          http://blog.sapien.com/current/2007/4/27" -ForegroundColor DarkGray
	Write-Host "          /creating-a-balloon-tip-notification-in-powershell.html" -ForegroundColor DarkGray
	Write-Host "          Code to extract icons from Shell32.dll by Thomas Levesque:"
	Write-Host "          http://stackoverflow.com/questions/6873026" -ForegroundColor DarkGray
	Write-Host
	Write-Host "Written by Rob van der Woude"
	Write-Host "http://www.robvanderwoude.com"
	Write-Host

	exit 1
}

$rc = 0

if ( $HOME[0] -eq '/' ) {
	# Linux: lshw/smartctl/df/notify-send commands

	# disk test requires root access, notification should NOT be run as root
	if ( ( . whoami ) -match "root" ) {
		Clear-Host
		Write-Host
		
		. smartctl -V > $null 2>&1
		if ( -not $? ) {
			Write-Host "This script requires smartmontools, available at"
			Write-Host "https://www.smartmontools.org/"
			exit 1
		}
	
		$notify = @( )
	
		Write-Host "Volume   `tStatus`tCapacity`tModel" -ForegroundColor White
		Write-Host "======   `t======`t========`t=====`n" -ForegroundColor White
	
		( . lshw -short -class disk ) -match "/dev/" | ForEach-Object {
			$disk = ( $_.trim( ) -split '\s+',  3 )[1]
			$name = ( $_.trim( ) -split 'disk', 2 )[1].trim( )
			if ( $disk -notmatch "(/cd|/dvd|/loop|/sr)" ) {
				try {
					$size = 0
					$size = ( ( ( . df -l --output=source,size ) -match $disk ) -split '\s+', 2 )[1] / 1MB
					if ( [int]$size -gt 0 ) {
						$test = ( ( ( . smartctl -H $disk ) -match "SMART [^\n\r]+: ([A-Z]+)" ) -split ":" )[1].trim( )
						if ( $test -eq "PASSED" ) {
								$fgc = "Green"
								$test = "OK"
						} else {
							$fgc = "Red"
							$rc = 1
						}
						$notify += "{0}\t{1}" -f $disk, $test
						Write-Host "$disk`t" -NoNewline
						Write-Host "$test`t" -ForegroundColor $fgc -NoNewline
						Write-Host ( "{0,5:N0} GB`t$name" -f $size ) -ForegroundColor White
					}
				}
				catch {
					# ignore errors from USB sticks etc.
				}
			}
		} 
	
		# prepare desktop notification:
		# a temporary shell script is used because notify-send
		# cannot send desktop notifications to other users
		if ( $rc -eq 1 ) {
			$icon = "error"
			$category = "device.error"
			$urgency = "critical"
			$title = "HDD Warnings"
		} else {
			$icon = "info"
			$category = "device"
			$urgency = "normal"
			$title = "HDD Status OK"
		}
		$message = $notify -join "\r"
		$notifycommand = "notify-send -i $icon -c $category -u $urgency `"$title`" `"$message`""
		# write notify-send command to temporary shell script
		Out-File -FilePath "$PSScriptRoot/GetHDDStatus.sh" -InputObject $notifycommand -Force
		Start-Sleep -s 1
		. chmod +x "$PSScriptRoot/GetHDDStatus.sh"
		Start-Sleep -s 1
	} else {
		if ( [System.IO.File]::Exists( "$PSScriptRoot/GetHDDStatus.sh" ) ) {
			Remove-Item -Path "$PSScriptRoot/GetHDDStatus.sh" -Force
		}
		# restart this PowerShell script as root, required for smartctl and lshw
		Start-Process -FilePath "sudo" -ArgumentList "pwsh `"$PSScriptRoot/GetHDDStatus.ps1`"" -NoNewWindow -Wait
		# run the root generated temporary shell script to show a desktop notification
		Start-Process -FilePath "bash" -ArgumentList "-c $PSScriptRoot/GetHDDStatus.sh"
		Start-Sleep -s 1
		Remove-Item -Path "$PSScriptRoot/GetHDDStatus.sh" -Force
	}
} else {
	# Windows: WMI and System.windows.Forms.NotifyIcon

	# based on code by Geoff @ UVM
	# https://www.uvm.edu/~gcd/2013/01/which-disk-is-that-volume-on/

	Clear-Host
	Write-Host
	
	[Reflection.Assembly]::LoadWithPartialName( "System.Windows.Forms" ) > $null 2>&1
	[Reflection.Assembly]::LoadWithPartialName( "System.Drawing" ) > $null 2>&1
	Add-Type -AssemblyName System.Windows.Forms

	[System.Collections.SortedList]$volumedetails = New-Object System.Collections.SortedList
	[System.Collections.SortedList]$volumestatus  = New-Object System.Collections.SortedList

	# query status for all local disk drives
	$diskdrives = Get-WmiObject -Namespace "root/CIMV2" -Class Win32_DiskDrive
	foreach ( $disk in $diskdrives ) {
		$diskindex  = $disk.Index
		$diskmodel  = $disk.Model
		$disksize   = "{0,5:F0} GB" -f ( $disk.Size / 1GB )
		$diskstatus = $disk.Status
		$part_query = 'ASSOCIATORS OF {Win32_DiskDrive.DeviceID="' + $disk.DeviceID.replace('\','\\') + '"} WHERE AssocClass=Win32_DiskDriveToDiskPartition'
		$partitions = @( Get-WmiObject -Query $part_query | Sort-Object StartingOffset )
		foreach ( $partition in $partitions ) {
			$vol_query = 'ASSOCIATORS OF {Win32_DiskPartition.DeviceID="' + $partition.DeviceID + '"} WHERE AssocClass=Win32_LogicalDiskToPartition'
			$volumes   = @( Get-WmiObject -Query $vol_query )
			foreach ( $volume in $volumes ) {
				# DriveType 3 means harddisks only
				# 0 = Unknown; 1 = No Root Directory; 2 = Removable Disk; 3 = Local Disk; 4 = Network Drive; 5 = Compact Disc; 6 = RAM Disk
				if ( $volume.DriveType -eq 3 ) {
					if ( -not $volumedetails.Contains( $volume.Name ) ) {
						$volumedetails.Add( $volume.Name, "[Disk $diskindex]  $disksize  $diskmodel" )
						$volumestatus.Add( $volume.Name, $diskstatus )
					}
				}
			}
		}
	}

	# console output table header
	Write-Host "Volume  Status  Disk#     Capacity  Model" -ForegroundColor White
	Write-Host "======  ======  =====     ========  =====`n" -ForegroundColor White

	# write table console output
	$volumedetails.Keys | ForEach-Object {
		$fgc = "Green"
		$status = ( $volumestatus[$_] )
		if ( $status -ne "OK" ) {
			$fgc = "Red"
			$rc  = 1
		}
		Write-Host ( "$_      " ) -ForegroundColor White -NoNewline
		Write-Host ( "$status      ".Substring( 0, 6 ) + "  " ) -ForegroundColor $fgc -NoNewline
		Write-Host ( $volumedetails[$_] ) -ForegroundColor White
	}

	# system tray balloon tip output
	[System.Windows.Forms.ToolTipIcon]$icon = [System.Windows.Forms.ToolTipIcon]::Info
	$title = "HDD Status OK"
	$systraymessage = ""
	$volumedetails.Keys | ForEach-Object {
		$status = ( $volumestatus[$_] )
		if ( $status -ne "OK" ) {
			$icon = [System.Windows.Forms.ToolTipIcon]::Error
			$title = "Warning: HDD Errors"
		}
		$systraymessage = $systraymessage + "$_`t$status`n"
	}

	# Extract system tray icon from Shell32.dll
	# C# code to extract icons from Shell32.dll by Thomas Levesque
	# http://stackoverflow.com/questions/6873026
	$signature = @'
	[DllImport( "Shell32.dll", EntryPoint = "ExtractIconExW", CharSet = CharSet.Unicode, ExactSpelling = true, CallingConvention = CallingConvention.StdCall )]
	private static extern int ExtractIconEx( string sFile, int iIndex, out IntPtr piLargeVersion, out IntPtr piSmallVersion, int amountIcons );

	public static Icon Extract( string file, int number, bool largeIcon )
	{
		IntPtr large;
		IntPtr small;
		ExtractIconEx( file, number, out large, out small, 1 );
		try
		{
			return Icon.FromHandle( largeIcon ? large : small );
		}
		catch
		{
			return null;
		}
	}
'@

	$iconextractor = Add-Type -MemberDefinition $signature -Name IconExtract -Namespace IconExtractor -ReferencedAssemblies System.Windows.Forms,System.Drawing -UsingNamespace System.Windows.Forms,System.Drawing -PassThru

	# icon depends on status
	if( $title -eq "HDD Status OK" ) {
		$systrayicon = $iconextractor::Extract( "C:\Windows\System32\shell32.dll", 223, $true )
	} else {
		$systrayicon = $iconextractor::Extract( "C:\Windows\System32\shell32.dll", 53, $true )
	}

	# show system tray icon and balloon tip
	$notify = New-Object System.windows.Forms.NotifyIcon
	$notify.BalloonTipText = $systraymessage
	$notify.BalloonTipTitle = $title
	$notify.BalloonTipIcon = [System.Windows.Forms.ToolTipIcon]$icon
	$notify.Icon = $systrayicon
	$notify.Visible = $true
	$notify.ShowBalloonTip( 30000 )
}

Write-Host

exit $rc