Added AssetHandler to allow for map_file_to_url. Fixes issue 153

git-svn-id: https://beef.googlecode.com/svn/trunk@802 b87d56ec-f9c0-11de-8c8a-61c5e9addfc9

lib/loader.rb

@@ -58,6 +58,8 @@ require 'lib/server/httpcontroller'
 
 require 'lib/server/httphookserver'
 
+require 'lib/server/assethandler'
+require 'lib/server/filehandler'
 require 'lib/server/zombiehandler'
 require 'lib/server/commandhandler'
 require 'lib/server/publichandler'

lib/modules/command.rb

@@ -94,7 +94,7 @@ module BeEF
     
     # Builds the 'datastore' attribute of the command which is used to generate javascript code.
     def build_datastore(data); 
-      @datastore = JSON.parse(data); 
+        @datastore = JSON.parse(data);
     end
     
     # Sets the datastore for the callback function. This function is meant to be called by the CommandHandler
@@ -280,6 +280,10 @@ module BeEF
       
       @beefjs_components[component] = component_path
     end
+
+    def map_file_to_url(file, path=nil, extension=nil, count=1)
+        return BeEF::AssetHandler.instance.bind(file, path, extension, count)
+    end
   
     private
     

lib/server/assethandler.rb

@@ -0,0 +1,74 @@
+module BeEF
+  
+  #
+  # Class defining BeEF assets 
+  #
+  class AssetHandler 
+    
+    # call BeEF::AssetHanlder.instance
+    include Singleton
+    
+    attr_reader :allocations, :root_dir
+    
+    def initialize
+      @allocations = {}
+      @http_server = BeEF::HttpHookServer.instance
+      @root_dir = File.expand_path('../../../', __FILE__)
+    end
+
+    #
+    # Binds a file to a mount point
+    #
+    def bind(file, path=nil, extension=nil, count=-1)
+        url = buildURL(path, extension) 
+        @allocations[url] = {'file' => "#{root_dir}"+file, 'path' => path, 'extension' => extension, 'count' => count} 
+        @http_server.mount(url, true, BeEF::FileHandler, @allocations[url]['file'])
+        puts @allocations
+        return url
+    end
+    
+    #
+    # Unbinds a file from a mount point
+    #
+    def unbind(url)
+        @allocations.delete(url)
+        @http_server.unmount(url, true)
+    end
+
+    #
+    # Builds a URL based on the path and extention, if neither are passed a random URL will be generated
+    #
+    def buildURL(path, extension, length=20)
+        url = (path == nil) ? '/'+rand(36**length).to_s(36) : path;
+        url += (extension == nil) ? '' : '.'+extension;
+        return url
+    end
+
+    #
+    # Checks if the file is allocated, if the file isn't return true to pass onto FileHandler.
+    #
+    def check(url)
+        if @allocations.has_key?(url)
+            count = @allocations[url]['count']
+            if count == -1
+                return true
+            end
+            if count > 0 
+                if (count - 1) == 0
+                    unbind(url)
+                else
+                    @allocations[url]['count'] = count - 1
+                end
+                return true
+            end
+        end
+        return false
+    end
+   
+   private
+   @http_server
+   @allocations
+
+  end
+  
+end

lib/server/filehandler.rb

@@ -0,0 +1,21 @@
+module BeEF
+  
+    #
+    # Custom FileHandler to deal with tracked files
+    #
+  class FileHandler < WEBrick::HTTPServlet::FileHandler
+  
+    #
+    # Override to do_GET to check tracked files
+    #
+    def do_GET(req, res)
+        if not BeEF::AssetHandler.instance.check(req.path)
+            raise HTPPStatus::NotFound, "`#req.path}` not found."
+        else 
+            super
+        end
+    end
+    
+  end
+  
+end

lib/server/httphookserver.rb

@@ -63,34 +63,34 @@ module BeEF
         config[:ServerSoftware] =  "BeEF " + VERSION
 
         @http_server = WEBrick::HTTPServer.new(config)
+        @asset_handler = BeEF::AssetHandler.instance
 
         # registers the ui pages
         Dir["#{$root_dir}/lib/ui/**/*.rb"].each { |http_module|
           require http_module
           mod_name = File.basename http_module, '.rb'
-          @http_server.mount "/ui/#{mod_name}", BeEF::HttpHandler, mod_name
+          mount("/ui/#{mod_name}", true, BeEF::HttpHandler, mod_name)
         }
         
         # registers the hook page
-        @http_server.mount "#{@configuration.get("hook_file")}", BeEF::ZombieHandler
-        @http_server.mount '/ui/public', BeEF::PublicHandler, "#{root_dir}/public"
-        @http_server.mount '/favicon.ico', WEBrick::HTTPServlet::FileHandler, "#{root_dir}#{@configuration.get("favicon_dir")}/#{@configuration.get("favicon_file_name")}"
-        @http_server.mount '/demos/', WEBrick::HTTPServlet::FileHandler, "#{root_dir}/demos/"
+        mount("#{@configuration.get("hook_file")}", true, BeEF::ZombieHandler)
+        mount('/ui/public', true, BeEF::PublicHandler, "#{root_dir}/public")
+        mount('/favicon.ico', true, WEBrick::HTTPServlet::FileHandler, "#{root_dir}#{@configuration.get("favicon_dir")}/#{@configuration.get("favicon_file_name")}")
+        mount('/demos/', true, WEBrick::HTTPServlet::FileHandler, "#{root_dir}/demos/")
 
         #dynamic handler
-        @http_server.mount '/dh', BeEF::DynamicHandler
+        mount('/dh', true, BeEF::DynamicHandler)
 
         #register mounts handled by dynamic handler
-        mounts['/init'] = BeEF::InitHandler
-        mounts['/event'] = BeEF::EventHandler
-        mounts['/requester'] = BeEF::RequesterHandler
+        mount('/init', false, BeEF::InitHandler)
+        mount('/event', false, BeEF::EventHandler)
+        mount('/requester', false, BeEF::RequesterHandler)
 
         # registers the command module pages
         Dir["#{root_dir}/modules/commands/**/*.rb"].each { |command|
           command_class = (File.basename command, '.rb').capitalize
           command_file = (File.basename command, '.rb')+'.js'
-          
-          mounts["/command/#{command_file}"] = BeEF::CommandHandler, command_class
+          mount("/command/#{command_file}", false, BeEF::CommandHandler, command_class)
         }
 
         trap("INT") { BeEF::HttpHookServer.instance.stop }
@@ -113,10 +113,40 @@ module BeEF
     # Restarts the BeEF http server.
     #
     def restart; stop; start; end
-    
+   
+    # 
+    # Mounts a handler, can either be a hard or soft mount (soft mounts are handled by the command handler
+    #
+    def mount(url, hard, c, args = nil)
+       if hard
+           if args == nil
+               @http_server.mount url, c
+            else
+               @http_server.mount url, c, *args
+            end
+        else
+            if args == nil
+                mounts[url] = c
+            else
+                mounts[url] = c, *args
+            end
+        end
+    end
+
+    #
+    # Unmounts handler
+    #
+    def unmount(url, hard)
+        if hard
+            @http_server.umount(url)
+        else
+            mounts.delete(url)
+        end
+    end
     
     private
     @http_server
+    @asset_handler
     
   end