PoshberryPi.psm1

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

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

    $LogicalDisk = Get-WmiObject Win32_LogicalDisk | Where-Object DeviceID -eq $TargetVolume
    $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 Wait-PiResponse {
    <#
        .SYNOPSIS
            Internal looping function which receives data after invoking Invoke-PiCommand.
 
        .DESCRIPTION
            Internal looping function which receives data after invoking Invoke-PiCommand.
 
        .PARAMETER TcpClient
            TCP Socket endpoint object used receive data
 
        .PARAMETER ServerStream
            Bytestream for sending and receiving data
 
        .NOTES
            Name: Wait-PiResponse
            Author: Boe Prox
            DateCreated: 22 Feb 2014
            Version History:
                Version 1.2 -- 4 Apr 2018
                    -Modified by Eli Hess to accomodate needs for dot net core
 
        .EXAMPLE
            Wait-PiResponse -TcpClient $TcpClient -ServerStream $ServerStream
 
            Description
            -----------
            Call typically made internally on Invoke-PiCommand utilizing TcpClient and Server stream
            already created for sending data.
    #>

[cmdletbinding()]
param (
    $TcpClient,
    $ServerStream
)
    $stringBuilder = New-Object Text.StringBuilder
    $Waiting = $True
    While ($Waiting) {
        While ($TcpClient.available -gt 0) {
            Write-Verbose "Processing return bytes: $($TcpClient.Available)"
            [byte[]]$inStream = New-Object byte[] $TcpClient.Available
            $buffSize = $TcpClient.Available
            $return = $ServerStream.Read($inStream, 0, $buffSize)
            [void]$stringBuilder.Append([System.Text.Encoding]::ASCII.GetString($inStream[0..($return-1)]))
            Start-Sleep -Seconds 1
        }
        If ($stringBuilder.length -gt 0) {
            $returnedData = [System.Management.Automation.PSSerializer]::DeSerialize($stringBuilder.ToString())
            Remove-Variable String -ErrorAction SilentlyContinue
            $Waiting = $False
        }
    }
    Write-Output $returnedData
}

function Send-PiResponse {
    <#
        .SYNOPSIS
            Internally used by Start-PiServer to send data back to caller.
 
        .DESCRIPTION
            Internally used by Start-PiServer to send data back to caller. Will initially
            attempt to serialize utilizing PSSerializer and then revert to
            ConvertTo-CliXml if an error occurs.
 
        .PARAMETER Response
            Response data to send.
 
        .NOTES
            Name: Send-PiResponse
            Author: Boe Prox
            DateCreated: 22 Feb 2014
            Version History:
                Version 1.2 -- 4 Apr 2018
                    -Modified by Eli Hess to accomodate needs for dot net core
 
        .EXAMPLE
            Send-PiResponse -Response
 
            Description
            -----------
            Internally used by Start-PiServer to send data back to caller.
    #>

[cmdletbinding()]
Param (
    $Response
)
    Try {
        Write-Verbose "Serializing data before sending using PSSerializer"
        $ErrorActionPreference = 'stop'
        $serialized = [System.Management.Automation.PSSerializer]::Serialize($Response)
    } Catch {
    Write-Verbose "Serializing data before sending using ConvertTo-CliXml"
        $serialized = $Response | ConvertTo-CliXml
    }
    $ErrorActionPreference = 'Continue'
    #Resend the Data back to the client
    $bytes  = [text.Encoding]::Ascii.GetBytes($serialized)
    #Send the data back to the client
    Write-Verbose "Echoing $($bytes.count) bytes to $remoteClient"
    $Stream.Write($bytes,0,$bytes.length)
    $Stream.Flush()
}

function ConvertTo-CliXml {
    <#
        .SYNOPSIS
            Serializes PSObjects into CliXml formatted string data.
 
        .DESCRIPTION
            Serializes PSObjects into CliXml formatted string data.
 
        .PARAMETER InputObject
            Object to serialize.
 
        .NOTES
            #Function borrowed from Joel Bennett (http://poshcode.org/4544)
            #Original Author Oisin Grehan (http://poshcode.org/1672)
 
        .EXAMPLE
            ConvertTo-CliXml -InputObject $Services
 
            Description
            -----------
            Serializes PSObjects into CliXml formatted string data.
    #>

[CmdletBinding()]
param(
    [Parameter(Position=0, Mandatory=$true, ValueFromPipeline=$true)]
    [ValidateNotNullOrEmpty()]
    [PSObject[]]$InputObject
)
    begin {
        $type = [PSObject].Assembly.GetType('System.Management.Automation.Serializer')
        $ctor = $type.GetConstructor('instance,nonpublic', $null, @([System.Xml.XmlWriter]), $null)
        $sw = New-Object System.IO.StringWriter
        $xw = New-Object System.Xml.XmlTextWriter $sw
        $serializer = $ctor.Invoke($xw)
    }
    process {
        try {
            [void]$type.InvokeMember("Serialize", "InvokeMethod,NonPublic,Instance", $null, $serializer, [object[]]@($InputObject))

        } catch {
            Write-Warning "Could not serialize $($InputObject.GetType()): $_"
        }
    }
    end {
        [void]$type.InvokeMember("Done", "InvokeMethod,NonPublic,Instance", $null, $serializer, @())
        $sw.ToString()
        $xw.Close()
        $sw.Dispose()
    }
}

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 {
            $DriveLetter = "$($DriveLetter.Substring(0,1)):"
            break
        }
        default {
            $DriveLetter = "$($DriveLetter.Substring(0,1)):"
        }
    }
    return $DriveLetter
}

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 Get-DiskAccess {
    <#
    .SYNOPSIS
        Returns a Win32DiskAccess object if validations pass
 
    .DESCRIPTION
        Returns a Win32DiskAccess object if validations pass
 
    .PARAMETER TargetVolume
        Volume of mounted drive to access
 
    .EXAMPLE
        $_diskAccess = Get-DiskAccess -TargetVolume "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]$TargetVolume
    )
    $_diskAccess = New-Object -TypeName "Posh.DiskWriter.Win32DiskAccess"
    #Lock logical drive
    $success = $_diskAccess.LockDrive($TargetVolume);
    Write-Verbose "Drive lock is $success"
    if (!$success)
    {
        Write-Error "Failed to lock drive"
        return $null
    }
    return $_diskAccess
}

Function Invoke-PiCommand {
    <#
        .SYNOPSIS
            Used to send PowerShell commands to a remote listener. Use this command with -Command Exit
            to shut down TCP Server.
 
        .DESCRIPTION
            Used to send PowerShell commands to a remote listener. Waits for a return response
            and presents data returned from remote system. Use this command with -Command Exit
            to shut down TCP Server.
 
        .PARAMETER Computername
            Computer to send command to
 
        .PARAMETER Port
            Remote port to target command on system running TCP Server
 
        .PARAMETER SourcePort
            Use a different source port for endpoint
 
        .PARAMETER Command
            Command to send to the TCP Server. Recommonded to be contained using single quotes if not
            using a variable containing the commands.
 
        .NOTES
            Name: Send-Command
            Author: Boe Prox
            DateCreated: 22 Feb 2014
            Version History:
                Version 1.2 -- 4 Apr 2018
                    -Modified by Eli Hess to accomodate needs for dot net core
                Version 1.1 -- 24 Feb 2014
                    -Added -ImpersonationLevel which will allow for a specific level of impersonation or no
                    impersonation at all.
                    -Broke out commonly used commands into Private functions (ConvertFrom-CliXml,Wait-Response)
                    -Changed SourePort default value to a randomized port in case command needs to run again to avoid
                    duplicate endpoint issues when source port is in a TIME_WAIT state
                Version 1.0 -- 22 Feb 2014
                    -Initial Version
 
        .EXAMPLE
            Invoke-PiCommand -Computername '192.168.1.40' -Port 2656 -Command 'Get-Process | Select -First 1'
 
            Description
            -----------
            Sends a Get-Process command to Server on port 2656 and returns the first process.
    #>

    [cmdletbinding()]
    Param (
        [parameter(ValueFromPipeline=$True,ValueFromPipelineByPropertyName=$True)]
        [string]$Computername = $env:COMPUTERNAME,
        [parameter()]
        [int]$Port = 1655,
        [parameter()]
        [int]$SourcePort=(Get-Random -Minimum 1500 -Maximum 16000),
        [parameter(Mandatory=$True)]
        [string]$Command = 'Exit'
    )
    Begin {
        Write-Verbose ("PSCommandPath $PSCommandPath")
        $PSBoundParameters.GetEnumerator() | ForEach-Object {
            Write-Verbose $_
        }
        Try {
            Write-Verbose "Creating Endpoint <$SourcePort> on $env:COMPUTERNAME"
            $Endpoint = new-object System.Net.IPEndpoint ([ipaddress]::any,$SourcePort)
            $TcpClient = [Net.Sockets.TCPClient]$endpoint
        } Catch {
            Write-Warning $_.Exception.Message
            Break
        }
    }
    Process {
        Try {
            Write-Verbose "Initiating connection to $Computername <$Port>"
            $TcpClient.Connect($Computername,$Port)
            $ServerStream = $TcpClient.GetStream()
            #Make the recieve buffer a little larger
            $TcpClient.ReceiveBufferSize = 1MB
            ##Client
            Try {
                Write-Verbose "Sending command"
                $data = [text.Encoding]::Ascii.GetBytes($Command)
                Write-Verbose "Sending $($data.count) bytes to $Computername <$port>"
                $ServerStream.Write($data,0,$data.length)
                $ServerStream.Flush()
                Wait-PiResponse -ServerStream $ServerStream -TcpClient $TcpClient
            } Catch {
                Write-Warning $_.Exception.Message
            }
        } Catch {
            Write-Warning $_.Exception.Message
        }
    }
    End {
        Write-Verbose 'Closing connection'
        If ($ServerStream) {$ServerStream.Dispose()}
        If ($TcpClient) {$TcpClient.Dispose()}
    }
}

function Start-PiServer {
    <#
        .SYNOPSIS
            Used to start a basic TCP server on your Raspberry Pi.
 
        .DESCRIPTION
            Used to start a basic TCP server on your Raspberry Pi.
 
        .PARAMETER Port
            Remote port to target command on system running TCP Server
 
        .NOTES
            Name: Start-PiServer
            Author: Boe Prox
            DateCreated: 22 Feb 2014
            Version History:
                Version 1.2 -- 4 Apr 2018
                    -Modified by Eli Hess to accomodate needs for dot net core
 
        .EXAMPLE
            Start-PiServer -Port 2656
 
            Description
            -----------
            Creates a TCP listener on port 2656 which echos output back to the source
    #>

[CmdletBinding()]
param(
    $Port=1655
)
    #Create the Listener port
    $Listener = New-Object System.Net.Sockets.TcpListener -ArgumentList $Port

    #Start the listener; opens up port for incoming connections
    $Listener.Start()
    Write-Verbose "Server started on port $Port"
    $Active = $True
    While ($Active) {
        $incomingClient = $Listener.AcceptTcpClient()
        $remoteClient = $incomingClient.client.RemoteEndPoint.Address.IPAddressToString
        Write-Verbose ("New connection from $remoteClient")
        #Let it buffer for a second
        Start-Sleep -Milliseconds 1000

        #Get the data stream from connected client
        $stream = $incomingClient.GetStream()
        #Validate default credentials
        Try {
            $activeConnection = $True
            $stringBuilder = New-Object Text.StringBuilder
            While ($incomingClient.Connected) {
                #Is there data available to process
                If ($Stream.DataAvailable) {
                    Do {
                        [byte[]]$byte = New-Object byte[] 1024
                        Write-Verbose "$($incomingClient.Available) Bytes available from $($remoteClient)"
                        $bytesReceived = $Stream.Read($byte, 0, $byte.Length)
                        If ($bytesReceived -gt 0) {
                            Write-Verbose "$bytesReceived Bytes received from $remoteClient"
                            [void]$stringBuilder.Append([text.Encoding]::Ascii.GetString($byte[0..($bytesReceived - 1)]))
                        } Else {
                            $activeConnection = $False
                            Break
                        }
                    } While ($Stream.DataAvailable)
                    $string = $stringBuilder.ToString()
                    If ($stringBuilder.Length -gt 0) {
                        If ($string -match '^(Quit|Exit)') {
                            Write-Verbose "Message received from $($remoteClient):`n$($stringBuilder.ToString())"
                            Write-Verbose 'Shutting down...'
                            $data = "Shutting down TCP Server on $Computername <$Port>"
                            Send-PiResponse -Response $data
                            $Active = $False
                            $Stream.Close()
                            $Listener.Stop()
                        } Else {
                            Write-Verbose "Message received from $($remoteClient):`n$string"
                            Try {
                                $ErrorActionPreference = 'Stop'
                                Write-Verbose "Running command"
                                $Data = [scriptblock]::Create($string).Invoke()
                            } Catch {
                                $Data = $_.Exception.Message
                            }
                            If (-Not $Data) {
                                $Data = 'No data to return!'
                            }
                            Send-PiResponse -Response $Data
                        }
                    } Else {
                        Send-PiResponse -Response 'No data'
                    }
                    Write-Verbose "Closing session to $remoteClient"
                    $incomingClient.Close()
                }
                Start-Sleep -Milliseconds 1000
            }
        } Catch {
            Write-Warning $_.Exception.Message
            Try {
                Send-PiResponse -Response $_ -ErrorAction Stop
            } Catch {
                Write-Warning $_.Exception.Message
            }
            $Stream.Dispose()
            $incomingClient.Close()
            $incomingClient.Dispose()
            Continue
        }
        [void]$stringBuilder.Clear()
    }
}

Function Write-PiImage {
    <#
    .SYNOPSIS
        Writes an image file to an SD card
 
    .DESCRIPTION
        Writes an image file to an SD card
 
    .PARAMETER TargetVolume
        Drive letter of mounted SD card
 
    .PARAMETER FileName
        Path to image file
 
    .EXAMPLE
        Write-PiImage -TargetVolume "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]$TargetVolume,
        [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
    }
    $TargetVolume = Format-DriveLetter $TargetVolume

    #Validate we're not targeting the system drive and the drive we're targeting is empty
    if($TargetVolume -eq $ENV:SystemDrive) {
        Write-Error "System Drive cannot be used as source"
        return $Completed
    } elseif ((Get-ChildItem $TargetVolume).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 -TargetVolume $TargetVolume
    }

    #Validate disk access is operational
    if($DiskAccess) {
        #Get drive size and open the physical drive
        $PhysicalDrive = Get-PhysicalDrive -TargetVolume $TargetVolume
        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 SourceVolume
        Drive letter of source SD card
 
    .PARAMETER FileName
        Full file path of img file to create
 
    .EXAMPLE
        Backup-PiImage -SourceVolume "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]$SourceVolume,
        [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
    $SourceVolume = Format-DriveLetter $SourceVolume
    #Validate we're not targeting the system drive
    if($SourceVolume -eq $ENV:SystemDrive) {
        Write-Error "System Drive cannot be targeted"
        return $Completed
    } else {
        $DiskAccess = Get-DiskAccess -TargetVolume $SourceVolume
    }

    if($DiskAccess) {
        #Get drive size and open the physical drive
        $PhysicalDrive = Get-PhysicalDrive -TargetVolume $SourceVolume
        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 TargetVolume
        Drive letter of target boot volume
 
    .EXAMPLE
        Enable-PiSSH -TargetVolume "D:"
 
        # Creates an empty file named 'ssh' on the boot volume mounted to D:
    .LINK
        https://github.com/eshess/PoshberryPi
 
    #>

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

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 TargetVolume
        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 -TargetVolume "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]$TargetVolume,
        [Parameter()]
        [string]$KeyMgmt = "WPA-PSK",
        [Parameter()]
        [System.Management.Automation.PSCredential]$WifiCredential,
        [Parameter()]
        [string]$CountryCode = "US",
        [Parameter()]
        [switch]$EncryptPSK
    )
    $TargetVolume = Format-DriveLetter $TargetVolume
    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 "$TargetVolume\wpa_supplicant.conf" -Encoding ascii
}