arachni 0.4.0.4 → 0.4.1

data/lib/arachni/rpc/server/base.rb

@@ -1,11 +1,17 @@
 =begin
-                  Arachni
-  Copyright (c) 2010-2012 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+    Copyright 2010-2012 Tasos Laskos <tasos.laskos@gmail.com>
 
-  This is free software; you can copy and distribute and modify
-  this program under the term of the GPL v2.0 License
-  (See LICENSE file for details)
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
 
+        http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
 =end
 
 require 'arachni/rpc/em'
@@ -17,10 +23,7 @@ class Server
 #
 # RPC server class
 #
-# @author: Tasos "Zapotek" Laskos
-#                                      <tasos.laskos@gmail.com>
-#                                      <zapotek@segfault.gr>
-# @version: 0.1
+# @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
 #
 class Base < ::Arachni::RPC::EM::Server
 

data/lib/arachni/rpc/server/dispatcher.rb

@@ -1,11 +1,17 @@
 =begin
-                  Arachni
-  Copyright (c) 2010-2012 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+    Copyright 2010-2012 Tasos Laskos <tasos.laskos@gmail.com>
 
-  This is free software; you can copy and distribute and modify
-  this program under the term of the GPL v2.0 License
-  (See LICENSE file for details)
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
 
+        http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
 =end
 
 require 'socket'
@@ -13,52 +19,51 @@ require 'sys/proctable'
 
 module Arachni
 
-require Options.instance.dir['lib'] + 'rpc/client/dispatcher'
-
-require Options.instance.dir['lib'] + 'rpc/server/base'
-require Options.instance.dir['lib'] + 'rpc/server/instance'
-require Options.instance.dir['lib'] + 'rpc/server/output'
+require Options.dir['lib'] + 'rpc/client'
+require Options.dir['lib'] + 'rpc/server/base'
+require Options.dir['lib'] + 'rpc/server/instance'
+require Options.dir['lib'] + 'rpc/server/output'
 
 module RPC
 class Server
 
 #
-# Dispatcher class
-#
-# Dispatches RPC servers on demand providing a centralized environment
-# for multiple RPC clients and allows for extensive process monitoring.
+# Dispatches RPC Instances on demand providing a centralized environment
+# for multiple clients and allows for extensive process monitoring.
 #
 # The process goes something like this:
-#   * a client issues a 'dispatch' call
-#   * the dispatcher starts a new RPC server on a random port
-#   * the dispatcher returns the port of the RPC server to the client
-#   * the client connects to the RPC server listening on that port and does his business
+# * on init the Dispatcher populates an Instance pool
+# * a client issues a 'dispatch' call
+# * the Dispatcher pops an Instance from the pool
+#   * asynchronously replenishes the pool
+#   * gives the Instance credentials to the client (url, auth token, etc.)
+# * the client connects to the Instance using these credentials
 #
-# Once the client finishes using the RPC server it *must* shut it down.<br/>
-# If it doesn't the system will be eaten away by idle instances of RPC servers.
+# Once the client finishes using the RPC Instance he *must* shut it down otherwise
+# the system will be eaten away by zombie RPC Instance processes.
 #
-# @author: Tasos "Zapotek" Laskos
-#                                      <tasos.laskos@gmail.com>
-#                                      <zapotek@segfault.gr>
-# @version: 0.2
+# @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
 #
 class Dispatcher
+    require Options.dir['lib'] + 'rpc/server/dispatcher/node'
+    require Options.dir['lib'] + 'rpc/server/dispatcher/handler'
 
-    require Options.instance.dir['lib'] + 'rpc/server/node'
-
-    include Arachni::Module::Utilities
-    include Arachni::UI::Output
+    include Utilities
+    include UI::Output
     include ::Sys
 
-    def initialize( opts )
 
+    HANDLER_LIB       = Options.dir['rpcd_handlers']
+    HANDLER_NAMESPACE = Handler
+
+    def initialize( opts )
         banner
 
         @opts = opts
 
-        @opts.rpc_port     ||= 7331
-        @opts.rpc_address  ||= 'localhost'
-        @opts.pool_size    ||= 5
+        @opts.rpc_port    ||= 7331
+        @opts.rpc_address ||= 'localhost'
+        @opts.pool_size   ||= 5
 
         if @opts.help
             print_help
@@ -66,12 +71,12 @@ class Dispatcher
         end
 
         @server = Base.new( @opts )
+        @server.logger.level = @opts.datastore[:log_level] if @opts.datastore[:log_level]
 
-        @server.add_async_check {
-            |method|
+        @server.add_async_check do |method|
             # methods that expect a block are async
-            method.parameters.flatten.include?( :block )
-        }
+            method.parameters.flatten.include? :block
+        end
 
         # let the instances in the pool know who to ask for routing instructions
         # when we're in grid mode.
@@ -79,37 +84,39 @@ class Dispatcher
 
         prep_logging
 
-        print_status( 'Initing RPC Server...' )
+        print_status 'Initing RPC Server...'
 
-        @server.add_handler( "dispatcher", self )
+        @server.add_handler( 'dispatcher', self )
 
         # trap interrupts and exit cleanly when required
         trap_interrupts { shutdown }
 
         @jobs = []
-        @pool = Queue.new
-        @replenisher = Queue.new
+        @consumed_pids = []
+        @pool = ::EM::Queue.new
 
-        print_status( 'Warming up the pool...' )
-        @opts.pool_size.times{ add_instance_to_pool }
-
-        # this thread will wait in the background and replenish the pool
-        Thread.new {
-            loop {
-                add_instance_to_pool
-                @replenisher.pop
-            }
-        }
+        if @opts.pool_size > 0
+            print_status 'Warming up the pool...'
+            @opts.pool_size.times{ add_instance_to_pool }
+        end
 
-        @node = nil
+        print_status 'Initialization complete.'
 
-        print_status( 'Done.' )
+        @node = Node.new( @opts, @logfile )
+        @server.add_handler( 'node', @node )
 
-        print_status( 'Initialization complete.' )
+        _handlers.each do |name, handler|
+            @server.add_handler( name, handler.new( @opts, self ) )
+        end
 
         run
     end
 
+    def handlers
+        _handlers.keys
+    end
+
+    # @return   [TrueClass]   true
     def alive?
         @server.alive?
     end
@@ -122,23 +129,27 @@ class Dispatcher
     #
     # @return   [Hash]      includes port number, owner, clock info and proc info
     #
-    def dispatch( owner = 'unknown', helpers = {} )
+    def dispatch( owner = 'unknown', helpers = {}, &block )
+        if @opts.pool_size <= 0
+            block.call false
+            return
+        end
 
         # just to make sure...
         owner = owner.to_s
-        cjob  = @pool.shift
-        cjob['owner']     = owner
-        cjob['starttime'] = Time.now
-        cjob['helpers']   = helpers
+        ::EM.next_tick { add_instance_to_pool }
+        @pool.pop do |cjob|
+            cjob['owner']     = owner
+            cjob['starttime'] = Time.now
+            cjob['helpers']   = helpers
 
-        print_status( "Instance dispatched -- PID: #{cjob['pid']} - " +
-            "Port: #{cjob['port']} - Owner: #{cjob['owner']}" )
+            print_status "Instance dispatched -- PID: #{cjob['pid']} - " +
+                "Port: #{cjob['port']} - Owner: #{cjob['owner']}"
 
-        @replenisher << true
+            @jobs << cjob
 
-        @jobs << cjob
-
-        return cjob
+            block.call cjob
+        end
     end
 
     #
@@ -149,18 +160,17 @@ class Dispatcher
     # @return   [Hash]
     #
     def job( pid )
-        @jobs.each {
-            |i|
-            cjob = i.dup
-            if cjob['pid'] == pid
-                cjob['currtime'] = Time.now
-                cjob['age'] = cjob['currtime'] - cjob['birthdate']
-                cjob['runtime']  = cjob['currtime'] - cjob['starttime']
-                cjob['proc'] =  proc( cjob['pid'] )
-
-                return cjob
-            end
-        }
+        @jobs.each do |j|
+            next if j['pid'] != pid
+            cjob = j.dup
+
+            cjob['currtime'] = Time.now
+            cjob['age']      = cjob['currtime'] - cjob['birthdate']
+            cjob['runtime']  = cjob['currtime'] - cjob['starttime']
+            cjob['proc']     = proc( cjob['pid'] )
+
+            return cjob
+        end
     end
 
     #
@@ -169,13 +179,7 @@ class Dispatcher
     # @return   [Array<Hash>]
     #
     def jobs
-        jobs = []
-        @jobs.each {
-            |cjob|
-            proc_info = job( cjob['pid'] )
-            jobs << proc_info if proc_info
-        }
-        return jobs
+        @jobs.map { |cjob| job( cjob['pid'] ) }.compact
     end
 
     #
@@ -185,55 +189,68 @@ class Dispatcher
     #
     def stats
         cjobs    = jobs( )
-        running  = cjobs.reject{ |job| job['proc'].empty? }
+        running  = cjobs.reject { |job| job['proc'].empty? }
         finished = cjobs - running
 
-        stats = {
-            'running_jobs'    => running,
-            'finished_jobs'   => finished,
-            'init_pool_size'  => @opts.pool_size,
-            'curr_pool_size'  => @pool.size
+        stats_h = {
+            'running_jobs'   => running,
+            'finished_jobs'  => finished,
+            'init_pool_size' => @opts.pool_size,
+            'curr_pool_size' => @pool.size,
+            'consumed_pids'  => @consumed_pids
         }
 
-        if @node
-            stats.merge!( 'node' => @node.info, 'neighbours' => @node.neighbours )
-        end
+        stats_h.merge!( 'node' => @node.info, 'neighbours' => @node.neighbours )
+
+        stats_h['node']['score']  = (rs_score = resource_consumption_score) > 0 ? rs_score : 1
+        stats_h['node']['score'] *= stats_h['node']['weight'] if stats_h['node']['weight']
+        stats_h['node']['score'] = Float( stats_h['node']['score'] )
 
-        return stats
+        stats_h
     end
 
+    # @return   [String]    contents of the log file
     def log
-        IO.read( prep_logging )
+        IO.read prep_logging
     end
 
+    # @return   [Hash]   the server's proc info
     def proc_info
-        p = proc( Process.pid )
+        proc( Process.pid ).merge( 'node' => @node.info )
+    end
 
-        if @node
-            p.merge!( 'node' => @node.info )
-        end
+    private
+
+    def self._handlers
+        @handlers ||= nil
+        return @handlers if @handlers
 
-        return p
+        @handlers = Component::Manager.new( HANDLER_LIB, HANDLER_NAMESPACE )
+        @handlers.load_all
+        @handlers
     end
 
-    private
+    def _handlers
+        self.class._handlers
+    end
+
+    def resource_consumption_score
+        cpu = mem = 0.0
+        jobs.each do |job|
+            mem += Float( job['proc']['pctmem'] ) if job['proc']['pctmem']
+            cpu += Float( job['proc']['pctcpu'] ) if job['proc']['pctcpu']
+        end
+        cpu + mem
+    end
 
     #
     # Outputs the Arachni banner.<br/>
     # Displays version number, revision number, author details etc.
     #
     def banner
-
-        puts 'Arachni - Web Application Security Scanner Framework
-       Author: Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
-                                      <zapotek@segfault.gr>
-               (With the support of the community and the Arachni Team.)
-
-       Website:       http://github.com/Zapotek/arachni
-       Documentation: http://github.com/Zapotek/arachni/wiki'
+        puts BANNER
         puts
         puts
-
     end
 
     def print_help
@@ -297,68 +314,58 @@ USAGE
 
 
     def trap_interrupts( &block )
-        [ 'QUIT', 'INT' ].each {
-            |signal|
+        %w(QUIT INT).each do |signal|
             trap( signal, &block || Proc.new{ } ) if Signal.list.has_key?( signal )
-        }
+        end
     end
 
     # Starts the dispatcher's server
     def run
-        print_status( 'Starting the server...' )
-        t = Thread.new { @server.run }
-
-        # wait for the server to settle
-        sleep( 0.1 ) while !@server.ready?
-
-        @node = Node.new( @opts, @logfile )
-        @server.add_handler( "node", @node )
-        t.join
+        print_status 'Starting the server...'
+        @server.start
     end
 
     def shutdown
-        print_status( 'Shutting down...' )
+        print_status 'Shutting down...'
         @server.shutdown
     end
 
     def add_instance_to_pool
-
         owner = 'dispatcher'
-        exception_jail{
+        exception_jail {
 
             # get an available port for the child
-            @opts.rpc_port = avail_port( )
-            @token         = secret( )
+            port  = avail_port
+            token = generate_token
 
             pid = ::EM.fork_reactor {
-                exception_jail {
-                    Arachni::RPC::Server::Instance.new( @opts, @token )
-                }
+                @opts.rpc_port = port
+                Server::Instance.new( @opts, token )
             }
 
-            print_status( "Instance added to pool -- PID: #{pid} - " +
-                "Port: #{@opts.rpc_port} - Owner: #{owner}" )
+            print_status "Instance added to pool -- PID: #{pid} - " +
+                "Port: #{@opts.rpc_port} - Owner: #{owner}"
 
             @pool << {
-                'token' => @token,
-                'pid'   => pid,
-                'port'  => @opts.rpc_port,
-                'url'   => "#{@opts.rpc_address}:#{@opts.rpc_port}",
-                'owner' => owner,
+                'token'     => token,
+                'pid'       => pid,
+                'port'      => port,
+                'url'       => "#{@opts.rpc_address}:#{port}",
+                'owner'     => owner,
                 'birthdate' => Time.now
             }
 
-            # let the child go about his business
+            @consumed_pids << pid
+
+            # let the child go about its business
             Process.detach( pid )
-            @token = nil
         }
-
     end
 
     def prep_logging
         # reroute all output to a logfile
         @logfile ||= reroute_to_file( @opts.dir['logs'] +
-            "Dispatcher - #{Process.pid}-#{@opts.rpc_port}.log" )
+            "/Dispatcher - #{Process.pid}-#{@opts.rpc_port}.log" )
     end
 
     def proc( pid )
@@ -367,16 +374,14 @@ USAGE
 
     def struct_to_h( struct )
         hash = {}
-
         return hash if !struct
 
-        struct.each_pair {
-            |k, v|
+        struct.each_pair do |k, v|
             v = v.to_s if v.is_a?( Bignum ) || v.is_a?( Fixnum )
             hash[k.to_s] = v
-        }
+        end
 
-        return hash
+        hash
     end
 
     #
@@ -385,13 +390,8 @@ USAGE
     # @return   Fixnum  port number
     #
     def avail_port
-
-        port = rand_port
-        while !avail_port?( port )
-            port = rand_port
-        end
-
-        return port
+        nil while !avail_port?( port = rand_port )
+        port
     end
 
     #
@@ -404,30 +404,27 @@ USAGE
         range[ rand( range.last - range.first ) ]
     end
 
-    def secret
-      secret = ''
-      1000.times {
-          secret += rand( 1000 ).to_s
-      }
-
-      return Digest::MD5.hexdigest( secret )
+    def generate_token
+        secret = ''
+        1000.times { secret << rand( 1000 ).to_s }
+        Digest::MD5.hexdigest( secret )
     end
 
     #
     # Checks whether the port number is available
     #
-    # @param    Fixnum  port
+    # @param    [Fixnum]  port
     #
-    # @return   Bool
+    # @return   [Bool]
     #
     def avail_port?( port )
         begin
             socket = Socket.new( :INET, :STREAM, 0 )
-            socket.bind( Addrinfo.tcp( "127.0.0.1", port ) )
+            socket.bind( Addrinfo.tcp( '127.0.0.1', port ) )
             socket.close
-            return true
+            true
         rescue
-            return false
+            false
         end
     end