Fix function parsing

Author: Meatballs1

lib/rex/exploitation/powershell/function.rb

@@ -37,12 +37,26 @@ def populate_params
       # 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

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
 

spec/lib/rex/exploitation/powershell/function_spec.rb

@@ -35,7 +35,12 @@
 
             [Parameter( Position = 1 )]
             [Type]
-            $ReturnType = [Void]
+            $ReturnType = [Void],
+
+            [String]$Parpy='hello',
+            [Integer] $puppy = 1,
+
+            [Array[]] $stuff = Array[],
         )
 
         $Domain = [AppDomain]::CurrentDomain
@@ -68,7 +73,11 @@
       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