#!/usr/bin/pwsh
#
# By MH 2020, MIT-License
# TODO: Logging (Sucess, Error), Windows-Support

param (
    [switch]$smart_disco, 
    [switch]$smart_check,
    [switch]$smart_with_external 
)


### Misc Helper Functions

function Get-IsUSBDrive {
    param (
        [ValidateNotNullOrEmpty()]        
        [string]$DeviceFile
    )

    $x = lsblk -d $DeviceFile -o SUBSYSTEMS |  Out-String -stream |  Select-String usb
    $x.Matches.Length -gt 0
}


### Smart-Helper Functions


class SmartDev {
    [String]$DeviceFile
    [String]$Name
    [String]$Type            
}


# First letter lower-case for compat with zabbix-expected Json
class SmartDiscoEntry {
    [String]$name
    [String]$type
    [String]$model
    [String]$sn
    [int]$rotations
    [int]$lbs
    [int]$isExternal
    [int]$isSSD
}


class SmartInfoEntry {
    [String]$Name
    [String]$DeviceFile
    [String]$Health
    [int]$PowerOnHours
    [int]$PowerCycleCount
    [double]$Temperature
}


function Get-SmartDevsFromScan {

    $Res=@()

    foreach ($ln in (smartctl -n standby --scan  | Out-String -stream)) {
        $Obj = New-Object SmartDev
        $parts=(-Split $ln)
        $Obj.DeviceFile=$parts[0]
        $Obj.Name=($parts[0] -Split "/")[2]
        $Obj.Type=(($ln -Split ",")[1] -Split " ")[1]
        $Res += $Obj
    }
    
    return $Res
}

function Get-SmartDisco {

    $Res = @()

    foreach ($dev in Get-SmartDevsFromScan) { 
        $isUSB=Get-IsUSBDrive -DeviceFile $dev.DeviceFile
                

        $model = ""  
        
        # TODO Support SSD
        $isSSD=0       
        if ( $isUSB -eq $False ) {
            $j = smartctl -n standby -a $dev.DeviceFile --json | ConvertFrom-Json
        
            $model = $j.model_name
            $sn = $j.serial_number
            $rotrate = $j.rotation_rate
            $lbs = $j.logical_block_size            
        }
        else {
            $model = $null
            $sn = $null
            $rotrate = $null
            $lbs = $null
        }
        
        $Obj=New-Object SmartDiscoEntry
        $Obj.name=$dev.Name
        $Obj.type=$dev.Type
        $Obj.model=$model
        $Obj.sn=$sn
        $Obj.rotations=$rotrate
        $Obj.lbs=$lbs
        $Obj.isExternal=$isUSB ? 1 : 0
        $Obj.isSSD=$isSSD
        $Res += $Obj
    }


    return $Res
}



function Get-SmartInfo {
    param (
        [Parameter(Mandatory=$true)]
        [SmartDev]$device
    )

    $Res = New-Object SmartInfoEntry
    $Res.DeviceFile = $device.DeviceFile
    $Res.Name = $device.Name

    smartctl -s on $device.DeviceFile | out-null
    
    $st=$(smartctl -n standby -q errorsonly -H $device.DeviceFile)

    if ( $st.Trim.Length -eq 0 ) {
        $st = "OK"
    }

    
    $j=smartctl -n standby -a $device.DeviceFile --json | ConvertFrom-Json

    $Res.Health = $st
    $Res.PowerOnHours = $j.power_on_time.hours
    $Res.PowerCycleCount = $j.power_cycle_count
    $Res.Temperature = $j.temperature.current
    
    return $Res
}


function Get-SmartInfos {

    param (
        [bool]$withExternalDevs
    )

    $Res = @()
    
    foreach ($dev in Get-SmartDevsFromScan) {
        $isUSB=Get-IsUSBDrive -DeviceFile $dev.DeviceFile

        if ( $isUSB -And -Not  $withExternalDevs ) {
            continue
        }
        
        smartctl -s on $dev.DeviceFile | out-null
        
        $si = Get-SmartInfo -device $dev
        $Res += $si
    }
        
    return $Res
}


### Monitoring functions

function Send-SmartDiscoToZabbix {
    $disco = Get-SmartDisco
    
    $jsons=$disco | ConvertTo-Json -Compress
    zabbix_sender -vv -k "8o_smartcheck.disco_devs" -o '$jsons'  -c /etc/zabbix/zabbix_agentd.conf
}


function Send-SmartInfosToZabbix {

    param (
        [bool]$withExternalDevs
    )

    &{
        foreach( $dev in Get-SmartInfos -withExternalDevs $withExternalDevs ) {
            # Better solution?
            $dn = $dev.Name
            Write-Output "- 8o_smartcheck.[${dn}.health] $($dev.Health)"
            Write-Output "- 8o_smartcheck.[${dn}.power_on_hours] $($dev.PowerOnHours)" 
            Write-Output "- 8o_smartcheck.[${dn}.power_cycle_count] $($dev.PowerCycleCount)" 
            Write-Output "- 8o_smartcheck.[${dn}.temperature] $($dev.Temperature)" 
        } 
    } | zabbix_sender -vv -i -  -c /etc/zabbix/zabbix_agentd.conf
}


$exitcode=0

if ( $smart_disco ) {
    Write-Output "Doing SMART-Disco "
    Send-SmartDiscoToZabbix
}

if ( $smart_check ) {
    Write-Output "Doing SMART-Check withExternal: ${smart_with_external} "
    Send-SmartInfosToZabbix -withExternalDevs $smart_with_external
}

Write-Output "Done"

exit $exitcode