New tree output for command wmap_sites -s with unicode option (wmap_s…

…ites -s [ids] (true/false))

Filtering 404 in tree output

plugins/wmap.rb

@@ -147,23 +147,28 @@ def cmd_wmap_sites(*args)
         when '-s'
           u = args.shift
           l = args.shift
-          s = args.shift
+          o = args.shift
 
-          if not u
-            return
-          end
+          return unless u
 
           if l == nil or l.empty?
             l = 200
-            s = true
+            o = 'true'
           else
-            l = l.to_i
-            s = false
+            # Add check if unicode parameters is the second one
+            if l == 'true' or l == 'false'
+              o = l
+              l = 200
+            else
+              l = l.to_i
+            end
           end
 
+          o = (o == 'true')
+
           if u.include? 'http'
             # Parameters are in url form
-            view_site_tree(u,l,s)
+            view_site_tree(u,l,o)
           else
               # Parameters are digits
               if !self.lastsites or self.lastsites.length == 0
@@ -188,12 +193,12 @@ def cmd_wmap_sites(*args)
             # Skip the DB entirely if no matches
             return if target_whitelist.length == 0
 
-            if not self.targets
-              self.targets = Hash.new()
-            end
+              unless self.targets
+                self.targets = Hash.new()
+              end
 
             target_whitelist.each do |ent|
-              view_site_tree(ent,l,s)
+              view_site_tree(ent,l,o)
             end
           end
           return
@@ -203,8 +208,7 @@ def cmd_wmap_sites(*args)
           print_line("\t-a [url]  Add site (vhost,url)")
           print_line("\t-d [ids]  Delete sites (separate ids with space)")
           print_line("\t-l        List all available sites")
-          print_line("\t-s [id]   Display site structure (vhost,url|ids) (level)")
-
+          print_line("\t-s [id]   Display site structure (vhost,url|ids) (level) (unicode output true/false)")
           print_line("")
           return
         else
@@ -1526,18 +1530,16 @@ def view_site_tree(urlstr, md, ld)
       # Skip the DB entirely if no matches
       return if site_whitelist.length == 0
 
-      vsites = Hash.new()
-
       site_whitelist.each do |ent|
         vhost,target = ent
 
         host = self.framework.db.workspace.hosts.find_by_address(target.host)
-        if not host
+        unless host
           print_error("No matching host for #{target.host}")
           next
         end
         serv = host.services.find_by_port_and_proto(target.port, 'tcp')
-        if not serv
+        unless serv
           print_error("No matching service for #{target.host}:#{target.port}")
           next
         end
@@ -1552,69 +1554,106 @@ def view_site_tree(urlstr, md, ld)
       end
     end
 
+    # Private function to avoid duplicate code
+    def load_tree_core(req, wtree)
+      pathchr = '/'
+      tarray = req.path.to_s.split(pathchr)
+      tarray.delete("")
+      tpath = Pathname.new(pathchr)
+      tarray.each do |df|
+        wtree.add_at_path(tpath.to_s,df)
+        tpath = tpath + Pathname.new(df.to_s)
+      end
+    end
+
     #
     # Load website structure into a tree
     #
-
     def load_tree(s)
-
-      pathchr = '/'
-
       wtree = Tree.new(s.vhost)
 
       # Load site pages
       s.web_pages.order('path asc').each do |req|
-        tarray = req.path.to_s.split(pathchr)
-        tarray.delete("")
-        tpath = Pathname.new(pathchr)
-        tarray.each do |df|
-          wtree.add_at_path(tpath.to_s,df)
-          tpath = tpath + Pathname.new(df.to_s)
+        if req.code != 404
+          load_tree_core(req, wtree)
         end
       end
 
       # Load site forms
       s.web_forms.each do |req|
-        tarray = req.path.to_s.split(pathchr)
-        tarray.delete("")
-        tpath = Pathname.new(pathchr)
-        tarray.each do |df|
-          wtree.add_at_path(tpath.to_s,df)
-          tpath = tpath + Pathname.new(df.to_s)
-        end
+        load_tree_core(req, wtree)
       end
 
-      return wtree
+      wtree
     end
 
-    #
-    # Print Tree structure. Still ugly
-    #
+    def print_file(filename)
+      ext = File.extname(filename)
+      if %w(.txt .md).include? ext
+        print '%bld%red'
+      elsif %w(.css .js).include? ext
+        print '%grn'
+      end
 
-    def print_tree(tree, ip, maxlevel, limitlevel)
-      initab = " " * 4
-      indent = 6
-      if  tree != nil and tree.depth <= maxlevel
-        print initab + (" " * indent * tree.depth)
-        if tree.depth > 0
-          print "|"+("-" * (indent-1))+"/"
-        end
-        if tree.depth >= 0
-          if tree.depth == 0
-            print "[#{tree.name}] (#{ip})\n"+initab+(" " * indent)+"\n"
+      print_line("#{ filename }%clr")
+    end
 
+    #
+    # Recursive function for printing the tree structure
+    #
+    def print_tree_recursive(tree, max_level, indent, prefix, is_last, unicode)
+      if  tree != nil and tree.depth <= max_level
+        print (' ' * indent)
+
+        # Prefix serve to print the superior hierarchy
+        prefix.each { |bool|
+          if unicode
+            print (bool ? ' ' : '│') + (' ' * 3)
           else
-            c = tree.children.count
-            if c > 0
-              print tree.name	+ " (" + c.to_s+")\n"
-            else
-              print tree.name	+ "\n"
-            end
+            print (bool ? ' ' : '|') + (' ' * 3)
           end
+        }
+        if unicode
+          # The last children is special
+          print (is_last ? '└' : '├') + ('─' * 2) + ' '
+        else
+          print (is_last ? '`' : '|') + ('-' * 2) + ' '
+        end
+
+        c = tree.children.count
+
+        if c > 0
+          print_line "%bld%blu#{ tree.name }%clr (#{ c.to_s })"
+        else
+          print_file tree.name
         end
 
-        tree.children.each_pair do |name,child|
-          print_tree(child,ip,maxlevel,limitlevel)
+        i = 1
+        new_prefix = prefix + [is_last]
+        tree.children.each_pair { |_,child|
+          is_last = !(i < c)
+          print_tree_recursive(child, max_level, indent, new_prefix, is_last, unicode)
+          i += 1
+        }
+      end
+    end
+
+    #
+    # Print Tree structure. Less ugly
+    # Modified by Jon P.
+    #
+    def print_tree(tree, ip, max_level, unicode)
+      indent = 4
+      if  tree != nil and tree.depth <= max_level
+        if tree.depth == 0
+          print_line "\n" + (' ' * indent) + "%cya[#{tree.name}] (#{ip})%clr"
+        end
+
+        i = 1
+        c = tree.children.count
+        tree.children.each_pair do |_,child|
+          print_tree_recursive(child, max_level, indent, [], !(i < c), unicode)
+          i += 1
         end
 
       end