PoshberryPi.psm1

Function Get-EncryptedPSK {
    <#
    .SYNOPSIS
        Generates the 32 byte encrypted hex string wpa_supplicant uses to connect to wifi
 
    .DESCRIPTION
        Generates the 32 byte encrypted hex string wpa_supplicant uses to connect to wifi
 
    .PARAMETER Credential
        A credential object containing the SSID and PSK used to connect to wifi
 
    .EXAMPLE
        $EncryptedPSK = Get-EncryptedPSK -Credential $Credential
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param (
        [Parameter()]
        [System.Management.Automation.PSCredential]$WifiCredential
    )
    if(!$PSBoundParameters.ContainsKey("WifiCredential"))
    {
        $WifiCredential = Get-Credential -Message "Please enter your Network SSID in the username field and passphrase as the password"
    }
    $NetCred = $WifiCredential.GetNetworkCredential()
    $Salt = [System.Text.Encoding]::ASCII.GetBytes($WifiCredential.UserName)
    $rfc = [System.Security.Cryptography.Rfc2898DeriveBytes]::New($NetCred.Password,$Salt,4096)
    Write-Output (Convert-ByteArrayToHexString -ByteArray $rfc.GetBytes(32) -Delimiter "").ToLower()
}

Function Get-IsPowerOfTwo {
    <#
    .SYNOPSIS
        Verifies input is a power of two and returns true or false
 
    .DESCRIPTION
        Verifies input is a power of two and returns true or false
 
    .PARAMETER Num
        Number to check against
 
    .EXAMPLE
        Get-IsPowerOfTwo -Num 23
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param (
        $Num
    )
    return ($Num -ne 0) -and (($Num -band ($Num - 1)) -eq 0);
}

Function Get-PhysicalDrive {
    <#
    .SYNOPSIS
        Returns the physical drive path of the DiskAccess object
 
    .DESCRIPTION
        Returns the physical drive path of the DiskAccess object
 
    .PARAMETER DriveLetter
        Volume to get physical path to
 
    .EXAMPLE
        $PhysicalDrive = Get-PhysicalDrive -DriveLetter "D:"
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [string]$DriveLetter
    )
    #Map to physical drive

    $LogicalDisk = Get-WmiObject Win32_LogicalDisk | Where-Object DeviceID -eq $DriveLetter
    $Log2Part = Get-WmiObject Win32_LogicalDiskToPartition | Where-Object Dependent -eq $LogicalDisk.__Path
    $phys = Get-WmiObject Win32_DiskDriveToDiskPartition | Where-Object Dependent -eq $Log2Part.Antecedent
    $DiskDrive = Get-WmiObject Win32_DiskDrive | Where-Object __Path -eq $phys.Antecedent
    Write-Verbose "Physical drive path is $($DiskDrive.DeviceID)"
    if($DiskDrive) {
        return $DiskDrive
    }else {
        Write-Error "Drive map unsuccessful"
        return $null
    }
}

Function Get-DiskHandle {
    <#
    .SYNOPSIS
        Opens the physical disk and returns the handle
 
    .DESCRIPTION
        Opens the physical disk and returns the handle
 
    .PARAMETER DiskAccess
        DiskAccess object to target
 
    .EXAMPLE
        $PhysicalHandle = Get-DiskHandle -DiskAccess $DiskAccess
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [Posh.DiskWriter.Win32DiskAccess]$DiskAccess,
        [parameter(Mandatory=$true)]
        [string]$PhysicalDrive
    )
    $physicalHandle = $DiskAccess.Open($PhysicalDrive)
    Write-Verbose "Physical handle is $physicalHandle"
    if ($physicalHandle -eq -1)
    {
        Write-Error "Failed to open physical drive"
        return $false
    }else {
        return $true
    }
}

Function Convert-ByteArrayToHexString {
    <#
    .SYNOPSIS
        Returns a hex representation of a System.Byte[] array as one or more strings. Hex format can be changed.
 
    .DESCRIPTION
        Returns a hex representation of a System.Byte[] array as one or more strings. Hex format can be changed.
 
    .PARAMETER ByteArray
        System.Byte[] array of bytes to put into the file. If you pipe this array in, you must pipe the [Ref] to the array.
        Also accepts a single Byte object instead of Byte[].
 
    .PARAMETER Width
        Number of hex characters per line of output.
 
    .PARAMETER Delimiter
        How each pair of hex characters (each byte of input) will be delimited from the next pair in the output. The default
        looks like "0x41,0xFF,0xB9" but you could specify "\x" if you want the output like "\x41\xFF\xB9" instead. You do
        not have to worry about an extra comma, semicolon, colon or tab appearing before each line of output. The default
        value is ",0x".
 
    .Parameter Prepend
        An optional string you can prepend to each line of hex output, perhaps like '$x += ' to paste into another
        script, hence the single quotes.
 
    .PARAMETER AddQuotes
        A switch which will enclose each line in double-quotes.
 
    .EXAMPLE
        [Byte[]] $x = 0x41,0x42,0x43,0x44
        Convert-ByteArrayToHexString $x
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
    [Parameter(Mandatory = $True, ValueFromPipeline = $True)]
    [System.Byte[]] $ByteArray,
    [Parameter()]
    [Int] $Width = 10,
    [Parameter()]
    [String] $Delimiter = ",0x",
    [Parameter()]
    [String] $Prepend = "",
    [Parameter()]
    [Switch] $AddQuotes
)
    if ($Width -lt 1)
    {
        $Width = 1
    }
    if ($ByteArray.Length -eq 0)
    {
        Write-Error "ByteArray length cannot be zero."
        Return
    }
    $FirstDelimiter = $Delimiter -Replace "^[\,\:\t]",""
    $From = 0
    $To = $Width - 1
    $Output = ""
    Do
    {
        $String = [System.BitConverter]::ToString($ByteArray[$From..$To])
        $String = $FirstDelimiter + ($String -replace "\-",$Delimiter)
        if ($AddQuotes)
        {
            $String = '"' + $String + '"'
        }
        if ($Prepend -ne "")
        {
            $String = $Prepend + $String
        }
        $Output += $String
        $From += $Width
        $To += $Width
    } While ($From -lt $ByteArray.Length)
    Write-Output $Output
}

Function Format-DriveLetter {
    <#
    .SYNOPSIS
        Returns uppercase driveletter with colon
 
    .DESCRIPTION
        Returns uppercase driveletter with colon
 
    .PARAMETER DriveLetter
        The string input to be validated
 
    .EXAMPLE
        $DriveLetter = Format-DriveLetter -DriveLetter "e"
 
        # Stores 'E:' in the variable DriveLetter
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [string]$DriveLetter
    )
    $DriveLetter = $DriveLetter.ToUpper()
    switch($DriveLetter.Length) {
        1 {
            $DriveLetter += ":"
            break
        }
        2 {
            break
        }
        default {
            $DriveLetter = "$($DriveLetter.Substring(0,1)):"
        }
    }
    return $DriveLetter
}

Function Get-DiskAccess {
    <#
    .SYNOPSIS
        Returns a Win32DiskAccess object if validations pass
 
    .DESCRIPTION
        Returns a Win32DiskAccess object if validations pass
 
    .PARAMETER DriveLetter
        Volume of mounted drive to access
 
    .EXAMPLE
        $_diskAccess = Get-DiskAccess -DriveLetter "D:"
 
        # Attempts to lock and open access to D: and return the access object to $_diskAccess
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [string]$DriveLetter
    )
    $_diskAccess = New-Object -TypeName "Posh.DiskWriter.Win32DiskAccess"
    #Lock logical drive
    $success = $_diskAccess.LockDrive($DriveLetter);
    Write-Verbose "Drive lock is $success"
    if (!$success)
    {
        Write-Error "Failed to lock drive"
        return $null
    }
    return $_diskAccess
}

Function Enable-PiWifi {
    <#
    .SYNOPSIS
        Enables wifi on the next boot of your Pi
 
    .DESCRIPTION
        Creates a 'wpa_supplicant.conf' file on the specified boot volume with desired settings to connect to wifi
 
    .PARAMETER KeyMgmt
        eg WPA-PSK
 
    .PARAMETER WifiCredential
        Credential object with the Username set to the WIFI SSID and the password set to the PSK
 
    .PARAMETER CountryCode
        eg US
 
    .PARAMETER Path
        Drive letter of boot volume
 
    .PARAMETER EncryptPSK
        Switch parameter for storing your PSK as encrypted text or plain text
 
    .EXAMPLE
        Enable-PiWifi -PSK $PSK -SSID $SSID -Path "D:"
 
        # Creates a 'wpa_supplicant.conf' file with default settings where possible on the boot volume mounted to D:
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param (
        [Parameter(Mandatory=$true)]
        [string]$Path,
        [Parameter()]
        [string]$KeyMgmt = "WPA-PSK",
        [Parameter()]
        [System.Management.Automation.PSCredential]$WifiCredential,
        [Parameter()]
        [string]$CountryCode = "US",
        [Parameter()]
        [switch]$EncryptPSK
    )
    if(!$PSBoundParameters.ContainsKey("WifiCredential"))
    {
        $WifiCredential = Get-Credential -Message "Please enter your Network SSID in the username field and passphrase as the password"
    }
    if($EncryptPSK){
        $PSK = Get-EncryptedPSK -WifiCredential $WifiCredential
    } else {
        $PSK = $WifiCredential.GetNetworkCredential().Password
    }
    $SSID = $WifiCredential.UserName
    $Output = @"
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1
country=$CountryCode
 
network={
    ssid="$SSID"
    psk=$PSK
    key_mgmt=$KeyMgmt
}
"@

    $Output.Replace("`r`n","`n") | Out-File "$Path\wpa_supplicant.conf" -Encoding ascii
}

Function Write-PiImage {
    <#
    .SYNOPSIS
        Writes an image file to an SD card
 
    .DESCRIPTION
        Writes an image file to an SD card
 
    .PARAMETER DriveLetter
        Drive letter of mounted SD card
 
    .PARAMETER FileName
        Path to image file
 
    .EXAMPLE
        Write-PiImage -DriveLetter "D:" -FileName "C:\Images\stretch.img"
 
        # Writes the image file located at C:\Images\stretch.img to the SD card mounted to D:
 
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param (
        [string]$DriveLetter,
        [string]$FileName
    )
    try { [Posh.DiskWriter.Win32DiskAccess] | Out-Null } catch { Add-Type -Path "$PSScriptRoot\classes\Win32DiskAccess.cs" }
    $Completed = $false
    $dtStart = (Get-Date)
    if((Test-Path $FileName) -eq $false)
    {
        Write-Error "$FileName doesn't exist"
        return $Completed
    }
    $DriveLetter = Format-DriveLetter $DriveLetter

    #Validate we're not targeting the system drive and the drive we're targeting is empty
    if($DriveLetter -eq $ENV:SystemDrive) {
        Write-Error "System Drive cannot be used as source"
        return $Completed
    } elseif ((Get-ChildItem $DriveLetter).Count -gt 0) {
        Write-Error "Target volume is not empty. Use diskpart to clean and reformat the target partition to FAT32."
        return $Completed
    } else {
        $DiskAccess = Get-DiskAccess -DriveLetter $DriveLetter
    }

    #Validate disk access is operational
    if($DiskAccess) {
        #Get drive size and open the physical drive
        $PhysicalDrive = Get-PhysicalDrive -DriveLetter $DriveLetter
        if($PhysicalDrive){
            $physicalHandle = Get-DiskHandle -DiskAccess $DiskAccess -PhysicalDrive $PhysicalDrive.DeviceID
        }
    }else {
        return $Completed
    }

    if($physicalHandle) {
        try {
            [console]::TreatControlCAsInput = $true
            $maxBufferSize = 1048576
            $buffer = [System.Array]::CreateInstance([Byte],$maxBufferSize)
            [long]$offset = 0;
            $fileLength = ([System.Io.FileInfo]::new($fileName)).Length
            $basefs = [System.Io.FileStream]::new($fileName, [System.Io.FileMode]::Open,[System.Io.FileAccess]::Read)
            $bufferOffset = 0;
            $BinanaryReader = [System.IO.BinaryReader]::new($basefs)
            while ($offset -lt $fileLength -and !$IsCancelling)
            {
                #Check for Ctrl-C and break if found
                if ([console]::KeyAvailable) {
                    $key = [system.console]::readkey($true)
                    if (($key.modifiers -band [consolemodifiers]"control") -and ($key.key -eq "C")) {
                        $IsCancelling = $true
                        break
                    }
                }

                [int]$readBytes = 0
                do
                {
                    $readBytes = $BinanaryReader.Read($buffer, $bufferOffset, $buffer.Length - $bufferOffset)
                    $bufferOffset += $readBytes
                } while ($bufferOffset -lt $maxBufferSize -and $readBytes -ne 0)

                [int]$wroteBytes = 0
                $bytesToWrite = $bufferOffset;
                $trailingBytes = 0;

                #Assume that the underlying physical drive will at least accept powers of two!
                if(Get-IsPowerOfTwo $bufferOffset)
                {
                    #Find highest bit (32-bit max)
                    $highBit = 31;
                    for (; (($bufferOffset -band (1 -shl $highBit)) -eq 0) -and $highBit -ge 0; $highBit--){}

                    #Work out trailing bytes after last power of two
                    $lastPowerOf2 = 1 -shl $highBit;

                    $bytesToWrite = $lastPowerOf2;
                    $trailingBytes = $bufferOffset - $lastPowerOf2;
                }

                if ($DiskAccess.Write($buffer, $bytesToWrite, [ref]$wroteBytes) -lt 0)
                {
                    Write-Error "Null disk handle"
                    return $Completed
                }

                if ($wroteBytes -ne $bytesToWrite)
                {
                    Write-Error "Error writing data to drive - past EOF?"
                    return $Completed
                }

                #Move trailing bytes up - Todo: Suboptimal
                if ($trailingBytes -gt 0)
                {
                    $Buffer.BlockCopy($buffer, $bufferOffset - $trailingBytes, $buffer, 0, $trailingBytes);
                    $bufferOffset = $trailingBytes;
                }
                else
                {
                    $bufferOffset = 0;
                }
                $offset += $wroteBytes;

                $percentDone = [int](100 * $offset / $fileLength);
                $tsElapsed = (Get-Date) - $dtStart
                $bytesPerSec = $offset / $tsElapsed.TotalSeconds;
                Write-Progress -Activity "Writing to Disk" -Status "Writing at $bytesPerSec" -PercentComplete $percentDone
            }
            $DiskAccess.Close()
            $DiskAccess.UnlockDrive()
            if(-not $IsCancelling) {
                $Completed = $true
                $tstotalTime = (Get-Date) - $dtStart
                Write-Verbose "All Done - Wrote $offset bytes. Elapsed time $($tstotalTime.ToString("dd\.hh\:mm\:ss"))"
            } else {
                Write-Output "Imaging was terminated early. Please clean and reformat the target volume before trying again."
            }
        } catch {
            $DiskAccess.Close()
            $DiskAccess.UnlockDrive()
        }finally {
            [console]::TreatControlCAsInput = $false
        }
    }
    return $Completed
}

Function Backup-PiImage {
    <#
    .SYNOPSIS
        Reads mounted SD card and saves contents to an img file
 
    .DESCRIPTION
        Reads mounted SD card and saves contents to an img file
 
    .PARAMETER DriveLetter
        Drive letter of source SD card
 
    .PARAMETER FileName
        Full file path of img file to create
 
    .EXAMPLE
        Backup-PiImage -DriveLetter "D:" -FileName "C:\Images\backup2018.img"
 
        # Creates a backup image of the SD card mounted to drive D: at C:\Images\backup2018.img
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [string]$DriveLetter,
        [parameter(Mandatory=$true)]
        [string]$FileName
    )
    try { [Posh.DiskWriter.Win32DiskAccess] | Out-Null } catch { Add-Type -Path "$PSScriptRoot\classes\Win32DiskAccess.cs" }
    $Completed = $false;
    $IsCancelling = $false
    $dtstart = Get-Date
    $maxBufferSize = 1048576
    $DriveLetter = Format-DriveLetter $DriveLetter
    #Validate we're not targeting the system drive
    if($DriveLetter -eq $ENV:SystemDrive) {
        Write-Error "System Drive cannot be targeted"
        return $Completed
    } else {
        $DiskAccess = Get-DiskAccess -DriveLetter $DriveLetter
    }

    if($DiskAccess) {
        #Get drive size and open the physical drive
        $PhysicalDrive = Get-PhysicalDrive -DriveLetter $DriveLetter
        if($PhysicalDrive){
            $readSize = $PhysicalDrive.Size
            $physicalHandle = Get-DiskHandle -DiskAccess $DiskAccess -PhysicalDrive $PhysicalDrive.DeviceID
        }
    }else {
        return $Completed
    }

    if($readSize -and $physicalHandle) {
        try {
            #Capture CTRL-C as input so we can free up disk locks
            [console]::TreatControlCAsInput = $true
            #Start doing the read
            $buffer =  [System.Array]::CreateInstance([Byte],$maxBufferSize)
            $offset = 0
            $fs = [System.Io.FileStream]::new($FileName, [System.Io.FileMode]::Create,[System.Io.FileAccess]::Write)
            while ($offset -lt $readSize -and !$IsCancelling)
            {
                #Check for CTRL-C and break if found
                if ([console]::KeyAvailable) {
                    $key = [system.console]::readkey($true)
                    if (($key.modifiers -band [consolemodifiers]"control") -and ($key.key -eq "C")) {
                        $IsCancelling = $true
                        break
                    }
                }
                #NOTE: If we provide a buffer that extends past the end of the physical device ReadFile() doesn't
                #seem to do a partial read. Deal with this by reading the remaining bytes at the end of the
                #drive if necessary
                if(($readSize - $offset) -lt $buffer.Length) {
                    $readMaxLength = $readSize - $offset
                } else {
                    $readMaxLength = $buffer.Length
                }
                [int]$readBytes = 0;
                if ($DiskAccess.Read($buffer, $readMaxLength, [ref]$readBytes) -lt 0)
                {
                    Write-Error "Error reading data from drive"
                    return $Completed;
                }
                if ($readBytes -eq 0)
                {
                    Write-Error "Error reading data from drive - past EOF?"
                    return $Completed
                }

                $fs.Write($buffer, 0, $readBytes)
                $offset += $readBytes

                $percentDone = (100*$offset/$readSize)
                $tsElapsed = (Get-Date) - $dtStart
                $bytesPerSec = $offset/$tsElapsed.TotalSeconds
                Write-Progress -Activity "Writing to disk" -Status "In Progress $bytesPerSec" -PercentComplete $percentDone
            }
            $fs.Close()
            $fs.Dispose()
            $DiskAccess.Close();
            $DiskAccess.UnlockDrive();
            $tstotalTime = (Get-Date) -$dtStart
        } catch {
            $DiskAccess.Close();
            $DiskAccess.UnlockDrive();
        } finally {
            [console]::TreatControlCAsInput = $false
        }
    }else {
        $DiskAccess.Close();
        $DiskAccess.UnlockDrive();
    }
    if (-not $IsCancelling)
    {
        $Completed = $true
        Write-Verbose "All Done - Read $offset bytes. Elapsed time $($tstotalTime.ToString("dd\.hh\:mm\:ss"))"
    }
    else
    {
        Write-Verbose "Cancelled";
        Remove-Item $FileName -Force
    }
    return $Completed
}

Function Enable-PiSSH {
    <#
    .SYNOPSIS
        Enables SSH remoting on next boot of your Pi
 
    .DESCRIPTION
        Creates an empty file named 'ssh' in the specified path. Placing this file in the boot volume of your Rasperry Pi
        will enable SSH remoting on next boot
 
    .PARAMETER Path
        Drive letter of target boot volume
 
    .EXAMPLE
        Enable-PiSSH -Path "D:"
 
        # Creates an empty file named 'ssh' on the boot volume mounted to D:
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

    [cmdletbinding()]
    param(
        [Parameter()]
        [String]$Path
    )
    New-Item -Path "$Path\" -Name ssh -ItemType File
}