diff --git a/core/modules.rb b/core/modules.rb index dac76bd1c..37d7ac953 100644 --- a/core/modules.rb +++ b/core/modules.rb @@ -33,11 +33,20 @@ module BeEF def self.get_categories categories = [] BeEF::Core::Configuration.instance.get('beef.module').each {|k,v| - if not categories.include?(v['category']) - categories << v['category'] + flatcategory = "" + if v['category'].kind_of?(Array) + # Therefore this module has nested categories (sub-folders), munge them together into a string with '/' characters, like a folder. + v['category'].each {|cat| + flatcategory << cat + "/" + } + else + flatcategory = v['category'] + end + if not categories.include?(flatcategory) + categories << flatcategory end } - return categories.sort + return categories.sort.uniq #This is now uniqued, because otherwise the recursive function to build the json tree breaks if there are duplicates. end # Get all modules currently stored in the database diff --git a/extensions/admin_ui/controllers/modules/modules.rb b/extensions/admin_ui/controllers/modules/modules.rb index 18528ceed..cfcf2ce2f 100644 --- a/extensions/admin_ui/controllers/modules/modules.rb +++ b/extensions/admin_ui/controllers/modules/modules.rb @@ -427,8 +427,28 @@ class Modules < BeEF::Extension::AdminUI::HttpController return BeEF::Module.support(mod, {'browser' => BD.get(hook_session_id, 'BrowserName'), 'ver' => BD.get(hook_session_id, 'BrowserVersion'), 'os' => [BD.get(hook_session_id, 'OsName')]}) end - def update_command_module_tree(tree, cmd_category, cmd_icon_path, cmd_status, cmd_name, cmd_id) + # If we're adding a leaf to the command tree, and it's in a subfolder, we need to recurse + # into the tree to find where it goes + def update_command_module_tree_recurse(tree,category,leaf) + working_category = category.shift + tree.each {|t| + if t['text'].eql? working_category and category.count > 0 + #We have deeper to go + update_command_module_tree_recurse(t['children'],category,leaf) + elsif t['text'].eql? working_category + #Bingo + t['children'].push(leaf) + break + end + } + + #return tree + + end + + #Add the command to the tree + def update_command_module_tree(tree, cmd_category, cmd_icon_path, cmd_status, cmd_name, cmd_id) # construct leaf node for the command module tree leaf_node = { 'text' => cmd_name, @@ -439,24 +459,67 @@ class Modules < BeEF::Extension::AdminUI::HttpController } # add the node to the branch in the command module tree - tree.each {|x| - if x['text'].eql? cmd_category - x['children'].push( leaf_node ) - break - end - } + if cmd_category.is_a?(Array) + #The category is an array, therefore it's a sub-folderised category + cat_copy = cmd_category.dup #Don't work with the original array, because, then it breaks shit + update_command_module_tree_recurse(tree,cat_copy,leaf_node) + else + #original logic here, simply add the command to the tree. + tree.each {|x| + if x['text'].eql? cmd_category + x['children'].push( leaf_node ) + break + end + } + end end - + + #Recursive function to build the tree now with sub-folders + def build_recursive_tree(parent,input) + cinput = input.shift.chomp('/') + if cinput.split('/').count == 1 #then we have a single folder now + if parent.detect {|p| p['text'] == cinput}.nil? + parent << {'text' => cinput, 'cls' => 'folder', 'children' => []} + else + if input.count > 0 + parent.each {|p| + if p['text'] == cinput + p['children'] = build_recursive_tree(p['children'],input) + end + } + end + end + else + #we have multiple folders + newinput = cinput.split('/') + newcinput = newinput.shift + if parent.detect {|p| p['text'] == newcinput }.nil? + fldr = {'text' => newcinput, 'cls' => 'folder', 'children' => []} + parent << build_recursive_tree(fldr['children'],newinput) + else + parent.each {|p| + if p['text'] == newcinput + p['children'] = build_recursive_tree(p['children'],newinput) + end + } + end + end + + if input.count > 0 + return build_recursive_tree(parent,input) + else + return parent + end + end + # Returns the list of all command_modules for a TreePanel in the interface. def select_command_modules_tree + blanktree = [] tree = [] - BeEF::Modules.get_categories.each { |c| - tree.push({ - 'text' => c, - 'cls' => 'folder', - 'children' => [] - }) - } + + #Due to the sub-folder nesting, we use some really badly hacked together recursion + #Note to the bored - if someone (anyone please) wants to refactor, I'll buy you cookies. -x + tree = build_recursive_tree(blanktree,BeEF::Modules.get_categories) BeEF::Modules.get_enabled.each{|k, mod| # get the hooked browser session id and set it in the command module @@ -510,7 +573,11 @@ class Modules < BeEF::Extension::AdminUI::HttpController # sort the children nodes by status and name tree.each {|x| x['children'] = - x['children'].sort_by {|a| [a['status'],a['text']]} + #Thanks a lot guys, thanks to the sub-folders this gets way more quirky .. + #Although, the sorting trick below is pretty funny :P + x['children'].sort_by {|a| + fldr = a['cls'] ? a['cls'] : 'zzzzz' + "#{fldr}#{a['status']}#{a['text']}"} } # append the number of command modules so the branch name results in: " (num)"