lib/ronin/ui/command_line/command_line.rb in ronin-0.2.2 vs lib/ronin/ui/command_line/command_line.rb in ronin-0.2.3

@@ -20,77 +20,97 @@
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 #++
 #
 
 require 'ronin/ui/command_line/exceptions/unknown_command'
-require 'ronin/ui/command_line/commands/console'
-require 'ronin/ui/console'
 
 require 'reverse_require'
+require 'extlib'
 
 module Ronin
   module UI
     module CommandLine
+      # Directory which stores the commands
+      COMMANDS_DIR = File.join('ronin','ui','command_line','commands')
+
+      # Name of the default to run
+      DEFAULT_COMMAND = 'console'
+
       #
       # Returns the commands registered with the command-line utility.
       #
       def CommandLine.commands
         unless class_variable_defined?('@@ronin_commands')
-          paths = Gem.find_resources('bin/ronin-*')
+          pattern = File.join('lib',COMMANDS_DIR,'*.rb')
+          paths = Gem.find_resources(pattern)
           
-          @@ronin_commands = {}
+          @@ronin_commands = []
             
           paths.each do |path|
-            next unless File.executable?(path)
-            name = File.basename(path).gsub(/^ronin-/,'')
+            name = File.basename(path).gsub(/\.rb$/,'')
 
-            @@ronin_commands[name] ||= path
+            @@ronin_commands << name unless @@ronin_commands.include?(name)
           end
         end
 
         return @@ronin_commands
       end
 
       #
-      # Returns +true+ if the a Command with the specified _name_ was
-      # registered with the command-line utility.
+      # Returns +true+ if a command exists with the specified _name_,
+      # returns +false+ otherwise.
       #
       def CommandLine.has_command?(name)
-        CommandLine.commands.has_key?(name.to_s)
+        CommandLine.commands.include?(name.to_s)
       end
 
       #
       # Returns the Command registered with the command-line utility
       # with the specified _name_.
       #
       def CommandLine.get_command(name)
         name = name.to_s
 
-        unless CommandLine.has_command?(name)
+        begin
+          require File.join(COMMANDS_DIR,name)
+        rescue LoadError
+          raise(UnknownCommand,"unable to load the command #{name.dump}",caller)
+        end
+
+        class_name = name.to_const_string
+
+        unless Commands.const_defined?(class_name)
           raise(UnknownCommand,"unknown command #{name.dump}",caller)
         end
 
-        return CommandLine.commands[name]
+        command = Commands.const_get(class_name)
+
+        unless command.respond_to?(:run)
+          raise(UnknownCommand,"command #{name.dump} must provide a 'run' method",caller)
+        end
+
+        return command
       end
 
       #
       # Runs the command-line utility with the given _argv_ Array. If the
       # first argument is a sub-command name, the command-line utility will
       # attempt to find and execute the Command with the same name.
       #
       def CommandLine.run(*argv)
-        if (argv.empty? || argv[0][0..0]=='-')
-          Commands::Console.run(*argv)
+        if (argv.empty? || argv.first[0..0]=='-')
+          name = DEFAULT_COMMAND
+          argv = ARGV
         else
-          cmd = argv.first
+          name = argv.first
           argv = argv[1..-1]
+        end
 
-          begin
-            exec(CommandLine.get_command(cmd),*argv)
-          rescue UnknownCommand => e
-            STDERR.puts e
-            exit -1
-          end
+        begin
+          CommandLine.get_command(name).run(*argv)
+        rescue UnknownCommand => e
+          STDERR.puts e
+          exit -1
         end
 
         return true
       end
     end