From dc382127414190c8ffbc4a173145dc307356252b Mon Sep 17 00:00:00 2001
From: Meatballs <eat_meatballs@hotmail.co.uk>
Date: Mon, 5 May 2014 20:53:36 +0100
Subject: [PATCH] Fix function parsing

---
 lib/rex/exploitation/powershell/function.rb   | 24 +++++++++++++++----
 lib/rex/exploitation/powershell/param.rb      |  2 +-
 .../exploitation/powershell/function_spec.rb  | 13 ++++++++--
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/lib/rex/exploitation/powershell/function.rb b/lib/rex/exploitation/powershell/function.rb
index 8468380bd0..76fcf5d86e 100644
--- a/lib/rex/exploitation/powershell/function.rb
+++ b/lib/rex/exploitation/powershell/function.rb
@@ -37,12 +37,26 @@ module Powershell
       # Get start of our block
       idx = scan_with_index('(',code[start..-1]).first.last + start
       pclause = block_extract(idx)
-      # Keep lines which declare a variable of some class
-      vars = pclause.split(/\n|;/).select {|e| e =~ /\]\$\w/}
-      vars.map! {|v| v.split('=',2).first}.map(&:strip)
+
+      func_regex = /\[(\w+\[\])\]\$(\w+)\s?=|\[(\w+)\]\$(\w+)\s?=|\[(\w+\[\])\]\s+?\$(\w+)\s+=|\[(\w+)\]\s+\$(\w+)\s?=/i
+      #func_regex = /\[(\w+\[\])\]\.?\$(\w+)\s?=|\[(\w+)\]\s?\$(\w+)\s?=/i
+      matches = pclause.scan(func_regex)
+
       # Ignore assignment, create params with class and variable names
-      vars.map {|e| e.split('$')}.each do |klass,name|
-        @params << Param.new(klass,name)
+      matches.each do |param|
+        klass = nil
+        name = nil
+        param.each do |value|
+          if value
+            if klass
+              name = value
+              @params << Param.new(klass,name)
+              break
+            else
+              klass = value
+            end
+          end
+        end
       end
     end
   end
diff --git a/lib/rex/exploitation/powershell/param.rb b/lib/rex/exploitation/powershell/param.rb
index 652f0af1ec..27815e9e1f 100644
--- a/lib/rex/exploitation/powershell/param.rb
+++ b/lib/rex/exploitation/powershell/param.rb
@@ -8,7 +8,7 @@ module Powershell
   class Param
     attr_accessor :klass, :name
     def initialize(klass,name)
-      @klass = klass.strip.gsub(/\[|\]|\s/,'')
+      @klass = klass.strip
       @name = name.strip.gsub(/\s|,/,'')
     end
 
diff --git a/spec/lib/rex/exploitation/powershell/function_spec.rb b/spec/lib/rex/exploitation/powershell/function_spec.rb
index 6ad49989b9..bb9159fe21 100644
--- a/spec/lib/rex/exploitation/powershell/function_spec.rb
+++ b/spec/lib/rex/exploitation/powershell/function_spec.rb
@@ -35,7 +35,12 @@ describe Rex::Exploitation::Powershell::Function do
 
             [Parameter( Position = 1 )]
             [Type]
-            $ReturnType = [Void]
+            $ReturnType = [Void],
+
+            [String]$Parpy='hello',
+            [Integer] $puppy = 1,
+
+            [Array[]] $stuff = Array[],
         )
 
         $Domain = [AppDomain]::CurrentDomain
@@ -68,7 +73,11 @@ describe Rex::Exploitation::Powershell::Function do
       function.code.should eq example_function_with_params
       function.to_s.include?("function #{function_name} #{example_function_with_params}").should be_true
       function.params.should be_kind_of Array
-      function.params.length.should be == 2
+      function.params.length.should be == 5
+      function.params[0].klass.should eq 'Type[]'
+      function.params[0].name.should eq 'Parameters'
+      function.params[1].klass.should eq 'Type'
+      function.params[1].name.should eq 'ReturnType'
     end
   end