Add-Persistence bugfix

When file paths were specified, they were not being properly validated.

2 files changed, 23 insertions, 11 deletions

diff --git a/Persistence/Persistence.psd1 b/Persistence/Persistence.psd1
index f51f7ce..e17faf1 100644
--- a/Persistence/Persistence.psd1
+++ b/Persistence/Persistence.psd1
@@ -4,7 +4,7 @@
 ModuleToProcess = 'Persistence.psm1'
 
 # Version number of this module.
-ModuleVersion = '1.1.0.0'
+ModuleVersion = '1.1.1.0'
 
 # ID used to uniquely identify this module
 GUID = '633d0f10-a056-41da-869d-6d2f75430195'
diff --git a/Persistence/Persistence.psm1 b/Persistence/Persistence.psm1
index 344c13d..7528f2e 100644
--- a/Persistence/Persistence.psm1
+++ b/Persistence/Persistence.psm1
@@ -417,11 +417,9 @@ function Add-Persistence
         [String]
         $PersistenceScriptName = 'Update-Windows',
 
-        [ValidateNotNullOrEmpty()]
         [String]
         $PersistentScriptFilePath = "$PWD\Persistence.ps1",
 
-        [ValidateNotNullOrEmpty()]
         [String]
         $RemovalScriptFilePath = "$PWD\RemovePersistence.ps1",
 
@@ -446,35 +444,49 @@ function Add-Persistence
         throw 'You provided invalid user-level persistence options.'
     }
 
-    $Path = Split-Path $PersistentScriptFilePath -ErrorAction Stop
+    $Result = Get-Item $PersistentScriptFilePath -ErrorAction SilentlyContinue
+    if ($Result -and $Result.PSIsContainer)
+    {
+        throw 'You must provide a file name with the PersistentScriptFilePath option.'
+    }
+
+    $Result = Get-Item $RemovalScriptFilePath -ErrorAction SilentlyContinue
+    if ($Result -and $Result.PSIsContainer)
+    {
+        throw 'You must provide a file name with the RemovalScriptFilePath option.'
+    }
+
+    $PersistentPath = Split-Path $PersistentScriptFilePath -ErrorAction Stop
     $Leaf = Split-Path $PersistentScriptFilePath -Leaf -ErrorAction Stop
     $PersistentScriptFile = ''
     $RemovalScriptFile = ''
 
-    if ($Path -eq '')
+    if ($PersistentPath -eq '')
     {
+        # i.e. Only a file name was provided implying $PWD
         $PersistentScriptFile = "$($PWD)\$($Leaf)"
     }
     else
     {
-        $PersistentScriptFile = "$($Path)\$($Leaf)"
+        $PersistentScriptFile = "$(Resolve-Path $PersistentPath)\$($Leaf)"
     }
 
-    $Path = Split-Path $RemovalScriptFilePath -ErrorAction Stop
+    $RemovalPath = Split-Path $RemovalScriptFilePath -ErrorAction Stop
     $Leaf = Split-Path $RemovalScriptFilePath -Leaf -ErrorAction Stop
-    if ($Path -eq '')
+    if ($RemovalPath -eq '')
     {
+        # i.e. Only a file name was provided implying $PWD
         $RemovalScriptFile = "$($PWD)\$($Leaf)"
     }
     else
     {
-        $RemovalScriptFile = "$($Path)\$($Leaf)"
+        $RemovalScriptFile = "$(Resolve-Path $RemovalPath)\$($Leaf)"
     }
 
     if ($PSBoundParameters['FilePath'])
     {
-        Get-ChildItem $FilePath -ErrorAction Stop
-        $Script = [IO.File]::ReadAllText((Resolve-Path $Path))
+        $null = Get-ChildItem $FilePath -ErrorAction Stop
+        $Script = [IO.File]::ReadAllText((Resolve-Path $FilePath))
     }
     else
     {