Land #4073, improved banner selection

lib/msf/base/config.rb

@@ -65,6 +65,7 @@ class Config < Hash
       'ModuleDirectory'     => "modules",
       'ScriptDirectory'     => "scripts",
       'LogDirectory'        => "logs",
+      'LogosDirectory'      => "logos",
       'SessionLogDirectory' => "logs/sessions",
       'PluginDirectory'     => "plugins",
       'DataDirectory'       => "data",
@@ -92,6 +93,13 @@ class Config < Hash
     self.new.config_directory
   end
 
+  # Return the directory that logo files should be loaded from.
+  #
+  # @return [String] path to the logos directory.
+  def self.logos_directory
+    self.new.logos_directory
+  end
+
   # Returns the global module directory.
   #
   # @return [String] path to global module directory.
@@ -148,6 +156,13 @@ class Config < Hash
     self.new.local_directory
   end
 
+  # Return the user-specific directory that logo files should be loaded from.
+  #
+  # @return [String] path to the logos directory.
+  def self.user_logos_directory
+    self.new.user_logos_directory
+  end
+
   # Returns the user-specific module base path
   #
   # @return [String] path to user-specific modules directory.
@@ -231,6 +246,13 @@ class Config < Hash
     InstallRoot
   end
 
+  # Return the directory that logo files should be loaded from.
+  #
+  # @return [String] path to the logos directory.
+  def logos_directory
+    data_directory + FileSep + self['LogosDirectory']
+  end
+
   # Returns the configuration directory default.
   #
   # @return [String] the root configuration directory.
@@ -301,6 +323,13 @@ class Config < Hash
     config_directory + FileSep + self['LocalDirectory']
   end
 
+  # Return the user-specific directory that logo files should be loaded from.
+  #
+  # @return [String] path to the logos directory.
+  def user_logos_directory
+    config_directory + FileSep + self['LogosDirectory']
+  end
+
   # Returns the user-specific module base path
   #
   # @return [String] path to user-specific modules directory.
@@ -339,6 +368,7 @@ class Config < Hash
     FileUtils.mkdir_p(session_log_directory)
     FileUtils.mkdir_p(loot_directory)
     FileUtils.mkdir_p(local_directory)
+    FileUtils.mkdir_p(user_logos_directory)
     FileUtils.mkdir_p(user_module_directory)
     FileUtils.mkdir_p(user_plugin_directory)
   end

lib/msf/ui/banner.rb

@@ -8,32 +8,23 @@ module Ui
 #
 ###
 module Banner
-
-  Logos =
-  %w{
-    branded-longhorn.txt
-    cow-head.txt
-    cowsay.txt
-    figlet.txt
-    i-heart-shells.txt
-    metasploit-shield.txt
-    missile-command.txt
-    ninja.txt
-    null-pointer-deref.txt
-    r7-metasploit.txt
-    wake-up-neo.txt
-    workflow.txt
-    3kom-superhack.txt
-    metasploit-park.txt
-  }
-
   #
-  # Returns a random metasploit logo.
+  # Returns a specific metasploit logo. If the specified file is a relative path
+  # then the file will be searched for first in the included local directory,
+  # then in the user-specific directory.
   #
   def self.readfile(fname)
-    base = File.expand_path(File.dirname(__FILE__))
-    pathname = File.join(base, "logos", fname)
-    fdata = "<< Missing banner: #{fname} >>"
+    pathname = fname
+
+    unless File.absolute_path(pathname) == pathname
+      if File.readable?(File.join(::Msf::Config.logos_directory, fname))
+        pathname = File.join(::Msf::Config.logos_directory, fname)
+      elsif File.readable?(File.join(::Msf::Config.user_logos_directory, fname))
+        pathname = File.join(::Msf::Config.user_logos_directory, fname)
+      end
+    end
+
+    fdata = "<< Missing banner: #{pathname} >>"
     begin
       raise ArgumentError unless File.readable?(pathname)
       raise ArgumentError unless File.stat(pathname).size < 4096
@@ -45,22 +36,19 @@ module Banner
   end
 
   def self.to_s
+    return self.readfile ENV['MSFLOGO'] if ENV['MSFLOGO']
+
+    logos = []
     # Easter egg (always a cow themed logo): export/set GOCOW=1
     if ENV['GOCOW']
-      case rand(3)
-        when 0
-          # branded-longhorn
-          self.readfile Logos[0]
-        when 1
-          # cow-head
-          self.readfile Logos[1]
-        else
-          # cowsay
-          self.readfile Logos[2]
-        end
+      logos.concat(Dir.glob(::Msf::Config.logos_directory + File::SEPARATOR + 'cow*.txt'))
     else
-      self.readfile Logos[rand(Logos.length)]
+      logos.concat(Dir.glob(::Msf::Config.logos_directory + File::SEPARATOR + '*.txt'))
+      logos.concat(Dir.glob(::Msf::Config.user_logos_directory + File::SEPARATOR + '*.txt'))
     end
+
+    logos = logos.map { |f| File.absolute_path(f) }
+    self.readfile logos[rand(logos.length)]
   end
 end
 

lib/msf/ui/logos/test.rb

@@ -1,6 +0,0 @@
-# -*- coding: binary -*-
-
-here = File.expand_path(File.dirname(__FILE__))
-
-puts "Hi I live #{here}!"
-