arachni 0.3 → 0.4

data/lib/arachni/ui/web/dispatcher_manager.rb

@@ -0,0 +1,274 @@
+=begin
+                  Arachni
+  Copyright (c) 2010-2012 Tasos "Zapotek" 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)
+
+=end
+
+require 'datamapper'
+require 'socket'
+
+require Arachni::Options.instance.dir['lib'] + 'rpc/client/dispatcher'
+
+module Arachni
+module UI
+module Web
+
+#
+# Provides methods for dispatcher management.
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1.1
+#
+class DispatcherManager
+
+    class Dispatcher
+        include DataMapper::Resource
+
+        property :id,           Serial
+        property :url,          String
+    end
+
+
+    def initialize( opts, settings )
+        @opts     = opts
+        @settings = settings
+
+        DataMapper::setup( :default, "sqlite3://#{@settings.db}/default.db" )
+        DataMapper.finalize
+
+        Dispatcher.auto_upgrade!
+    end
+
+    #
+    # Puts a new dispatcher (and it's neighbours) in the DB.
+    #
+    # @param    [String]    url          URL of the dispatcher
+    # @param    [Bool]      neighbours   add its neighbouring dispatchers too?
+    #
+    def new( url, neighbours = true )
+        Dispatcher.first_or_create( :url => url )
+
+        return if !neighbours
+        connect( url ).node.neighbours {
+            |neighbours|
+            neighbours.each {
+                |node|
+                Dispatcher.first_or_create( :url => node )
+            }
+        }
+    end
+
+    #
+    # Provides an easy way to connect to a dispatcher.
+    #
+    # @param    [String]   url
+    #
+    # @return   [Arachni::RPC::Client::Dispatcher]
+    #
+    def connect( url )
+        Arachni::RPC::Client::Dispatcher.new( @opts, url )
+    end
+
+    #
+    # Checks wether the dispatcher is alive.
+    #
+    # @param    [String]    url     URL of the dispatcher
+    #
+    def alive?( url, &block )
+        raise( "This method requires a block!" ) if !block_given?
+
+        connect( url ).alive? {
+            |ret|
+            block.call( ret.rpc_connection_error? ? false : true )
+        }
+    end
+
+    def first_alive( &block )
+        raise( "This method requires a block!" ) if !block_given?
+
+        if !all.empty?
+            EM.synchrony do
+                dispatchers = EM::Synchrony::Iterator.new( all ).map {
+                    |dispatcher, iter|
+                    alive?( dispatcher.url ){
+                        |bool|
+                        if bool
+                            iter.return( dispatcher )
+                        else
+                            iter.return( nil )
+                        end
+                    }
+                }.compact
+
+                if dispatchers.empty?
+                    block.call( false )
+                else
+                    block.call( dispatchers.pop )
+                end
+            end
+        else
+            block.call( false )
+        end
+    end
+
+    def jobs( &block )
+        ::EM::Iterator.new( all, 20 ).map( proc {
+            |dispatcher, iter|
+
+            connect( dispatcher.url ).stats {
+                |stats|
+                iter.return( stats['running_jobs'] ) if !stats.rpc_connection_error?
+            }
+
+        }, proc {
+            |running|
+            block.call( running.flatten )
+        })
+    end
+
+    #
+    # Provides statistics about running jobs etc using the dispatcher
+    #
+    # @return   [Hash]
+    #
+    def stats( &block )
+        raise( "This method requires a block!" ) if !block_given?
+
+        EM.synchrony do
+
+            stats = EM::Synchrony::Iterator.new( all ).map {
+                |dispatcher, iter|
+
+                if !dispatcher.rpc_connection_error?
+                    connect( dispatcher.url ).stats {
+                        |stats|
+                        if !stats.rpc_exception?
+
+                            # automatically grab and save neighbours
+                            stats['neighbours'].each { |n| new( n ) }
+
+                            iter.return( { dispatcher.url => stats } )
+                        else
+                            iter.return( nil )
+                        end
+                    }
+                end
+            }.compact
+
+            sorted_stats = {}
+            stats.sort{ |a, b| a.keys[0] <=> b.keys[0] }.each {
+                |stat|
+                sorted_stats.merge!( stat )
+            }
+
+            sorted_stats.each_pair {
+                |k, stats|
+
+                sorted_stats[k]['running_jobs'] =
+                    EM::Synchrony::Iterator.new( stats['running_jobs'] ).map {
+                    |instance, iter|
+
+                    if instance['helpers']['rank'] != 'slave'
+                        @settings.instances.connect( instance['url'] ).framework.progress_data(
+                            :slaves   => false,
+                            :messages => false,
+                            :issues   => false
+                        ) {
+                            |prog_data|
+                            if prog_data.rpc_exception?
+                                iter.return( nil )
+                            else
+                                instance.merge!( prog_data['stats'] )
+                                instance['status']  = prog_data['status'].capitalize!
+                                iter.return( instance )
+                            end
+                        }
+                    else
+                        @settings.instances.connect( instance['helpers']['master'] ).framework.progress_data(
+                            :messages => false,
+                            :issues   => false
+                        ) {
+                            |prog_data|
+                            if prog_data.rpc_exception?
+                                iter.return( nil )
+                            else
+                                prog_data['instances'].each {
+                                    |insdat|
+                                     if insdat['url'] == instance['url']
+                                         instance.merge!( insdat )
+                                         instance['status'].capitalize!
+                                         iter.return( instance )
+                                    end
+                                }
+                            end
+                        }
+                    end
+                }.compact
+            }
+
+            block.call( sorted_stats )
+        end
+    end
+
+    #
+    # Returns all dispatchers stored in the DB.
+    #
+    # @return    [Array]
+    #
+    def all( *args )
+        Dispatcher.all( *args )
+    end
+
+    def all_with_liveness( &block )
+        raise( "This method requires a block!" ) if !block_given?
+
+        EM::Iterator.new( all ).map( proc{
+            |dispatcher, iter|
+
+            alive?( dispatcher.url ) {
+                |liveness|
+
+                m_dispatcher = {}
+                dispatcher.attributes.each_pair {
+                    |k, v|
+                    m_dispatcher[k.to_s] = v
+                }
+                m_dispatcher['alive'] = liveness
+                iter.return( m_dispatcher )
+            }
+        }, proc{
+            |dispatchers|
+            block.call( dispatchers )
+        })
+    end
+
+    #
+    # Removed all dispatchers from the DB.
+    #
+    def delete_all
+        all.each {
+            |report|
+            delete( report.id )
+        }
+        all.destroy
+    end
+
+    #
+    # Removed a dispatcher from the DB.
+    #
+    # @param    [Integer]   id
+    #
+    def delete( id )
+        Dispatcher.get( id ).destroy
+    end
+
+end
+end
+end
+end

data/lib/arachni/ui/web/instance_manager.rb

@@ -0,0 +1,69 @@
+=begin
+                  Arachni
+  Copyright (c) 2010-2012 Tasos "Zapotek" 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)
+
+=end
+
+require Arachni::Options.instance.dir['lib'] + 'rpc/client/instance'
+require Arachni::Options.instance.dir['lib'] + 'ui/web/utilities'
+require Arachni::Options.instance.dir['lib'] + 'rpc/client/instance'
+
+module Arachni
+module UI
+module Web
+
+#
+# Provides methods for instance management.
+#
+# @author: Tasos "Zapotek" Laskos
+#                                      <tasos.laskos@gmail.com>
+#                                      <zapotek@segfault.gr>
+# @version: 0.1
+#
+class InstanceManager
+
+    include Utilities
+
+    def initialize( opts, settings )
+        @opts     = opts
+        @settings = settings
+    end
+
+    #
+    # Provides an easy way to connect to an instance
+    #
+    # @param    [String]   url
+    # @param    [Hash]     session  session of the current user (optional)
+    # @param    [String]   token    authentication token (optional)
+    #
+    # @return   [Arachni::RPC::Client::Instance]
+    #
+    def connect( url, session = nil, token = nil )
+        #
+        # Sync up the session authentication tokens with the ones in the
+        # class variables.
+        #
+        # This will allow users to still connect to instances even if they
+        # shutdown the WebUI or remove their cookies.
+        #
+        @@tokens ||= {}
+        session['tokens'] ||= {} if session
+        @@tokens[url] = token if token
+
+        session['tokens'].merge!( @@tokens ) if session
+        @@tokens.merge!( session['tokens'] ) if session
+        session['tokens'].merge!( @@tokens ) if session
+
+        tmp_token = session ? session['tokens'][url] : @@tokens[url]
+
+        Arachni::RPC::Client::Instance.new( @opts, url, tmp_token )
+    end
+
+end
+end
+end
+end

data/lib/{ui → arachni/ui}/web/log.rb

@@ -1,6 +1,6 @@
 =begin
                   Arachni
-  Copyright (c) 2010-2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+  Copyright (c) 2010-2012 Tasos "Zapotek" 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

data/lib/arachni/ui/web/output_stream.rb

@@ -0,0 +1,54 @@
+=begin
+                  Arachni
+  Copyright (c) 2010-2012 Tasos "Zapotek" 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)
+
+=end
+
+module Arachni
+module UI
+module Web
+
+    #
+    # This used to be a stream in the past, now it's just a compat class.
+    #
+    class OutputStream
+
+        def initialize( output, lines, &block )
+            @lines  = lines
+            @output = output
+            @buffer = []
+
+            @icon_whitelist = {}
+            [ 'status', 'ok', 'error', 'info', 'bad' ].each {
+                |icon|
+                @icon_whitelist[icon] = "<img src='/icons/#{icon}.png' />"
+            }
+
+        end
+
+        def format
+            str = ''
+            cnt = 0
+
+            while( ( out = @output.pop ) && ( ( cnt += 1 ) < @lines ) )
+
+                type = out.keys[0]
+                msg  = out.values[0]
+
+                next if out.values[0].empty?
+
+                icon = @icon_whitelist[type.to_s] || ''
+                str += icon + CGI.escapeHTML( " #{out.values[0]}" ) + "<br/>"
+            end
+
+            return str
+        end
+
+    end
+end
+end
+end

data/lib/{ui → arachni/ui}/web/report_manager.rb

@@ -1,6 +1,6 @@
 =begin
                   Arachni
-  Copyright (c) 2010-2011 Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
+  Copyright (c) 2010-2012 Tasos "Zapotek" 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
@@ -22,7 +22,7 @@ module Web
 # @author: Tasos "Zapotek" Laskos
 #                                      <tasos.laskos@gmail.com>
 #                                      <zapotek@segfault.gr>
-# @version: 0.1.1
+# @version: 0.2
 #
 class ReportManager
 
@@ -33,13 +33,12 @@ class ReportManager
         include DataMapper::Resource
 
         property :id,           Serial
-        property :host,         String
+        property :host,         Text
         property :issue_count,  Integer
-        property :filename,     String
+        property :filename,     Text
         property :datestamp,    DateTime
     end
 
-
     def initialize( opts, settings )
         @opts     = opts
         @settings = settings
@@ -48,6 +47,7 @@ class ReportManager
         DataMapper::setup( :default, "sqlite3://#{@settings.db}/default.db" )
         DataMapper.finalize
 
+        # Report.raise_on_save_failure = true
         Report.auto_upgrade!
 
         migrate_files
@@ -63,14 +63,17 @@ class ReportManager
             next if Report.first( :filename => File.basename( file, EXTENSION ) )
 
             begin
-                data = File.read( file )
+                report = ::Arachni::AuditStore.load( file )
                 Report.create(
-                    :issue_count => get_issue_count( data ),
-                    :host        => get_host( data ),
+                    :issue_count => get_issue_count( report ),
+                    :host        => get_host( report ),
                     :filename    => File.basename( file, EXTENSION ),
-                    :datestamp   => get_finish_datetime( data )
+                    :datestamp   => get_finish_datetime( report )
                 )
-            rescue
+            rescue Exception => e
+                # p file
+                # ap e
+                # ap e.backtrace
             end
         }
     end
@@ -79,7 +82,7 @@ class ReportManager
     # @return    [String]    save directory
     #
     def savedir
-        @settings.public + "/#{FOLDERNAME}/"
+        @settings.public_folder + "/#{FOLDERNAME}/"
     end
 
     #
@@ -92,25 +95,23 @@ class ReportManager
     #
     # Saves the report to a file
     #
-    # @param    [String]        report  YAML serialized audistore object as returned by the Arachni XMLRPC server.
-    #                                       Basically an 'afr' report as a string.
+    # @param    [Arachni::AuditStore]    report
     #
     # @return   [String]        the path to the saved report
     #
     def save( report )
-        @settings.log.report_saved( {}, get_filename( report ) )
+        @settings.log.report_saved( {}, report_to_filename( report ) )
         return save_to_file( report, report_to_path( report ) )
     end
 
     #
     # Gets the path to a given report based on the contents of the report
     #
-    # @param    [String]        report  YAML serialized audistore object as returned by the Arachni XMLRPC server.
-    #                                       Basically an 'afr' report as a string.
+    # @param    [Arachni::AuditStore]   report
     # @return   [String]
     #
     def report_to_path( report )
-        savedir + File.basename( get_filename( report ) + EXTENSION )
+        savedir + File.basename( report_to_filename( report ) + EXTENSION )
     end
 
     #
@@ -158,26 +159,25 @@ class ReportManager
     # Generates a filename based on the contents of the report in the form of
     # host:audit_date
     #
-    # @param    [String]        report  YAML serialized audistore object as returned by the Arachni XMLRPC server.
-    #                                       Basically an 'afr' report as a string.
+    # @param    [Arachni::AuditStore]    report
     #
-    # @return   [String]        host:audit_date
+    # @return   [String]        host.audit_date.ext
     #
-    def get_filename( report )
-        rep = unserialize( report )
-        filename = "#{URI(rep.options['url']).host}:#{rep.start_datetime}"
+    def report_to_filename( report )
+        filename = "#{URI(report.options['url']).host}:#{report.start_datetime}"
+        filename.gsub( ':', '.' ).gsub( ' ', '_' ).gsub( '-', '_' ).gsub( '__', '_' )
     end
 
     def get_issue_count( report )
-        unserialize( report ).issues.size
+        report.issues.size
     end
 
     def get_host( report )
-        return URI(unserialize( report ).options['url']).host
+        return URI( report.options['url'] ).host
     end
 
     def get_finish_datetime( report )
-        return unserialize( report ).finish_datetime
+        return report.finish_datetime
     end
 
     #
@@ -193,8 +193,16 @@ class ReportManager
 
         begin
             location = savedir + Report.get( id ).filename + EXTENSION
-            convert( type, File.read( location ) )
-        rescue
+
+            # if it's the default report type don't waste time converting
+            if '.' + type == EXTENSION
+                return File.read( location )
+            else
+                return convert( type, ::Arachni::AuditStore.load( location ) )
+            end
+        rescue Exception => e
+            ap e
+            ap e.backtrace
             return nil
         end
     end
@@ -219,50 +227,37 @@ class ReportManager
 
     private
 
-    def unserialize( data )
-         begin
-            Marshal.load( data )
-         rescue
-             YAML.load( data )
-         end
-    end
-
-    def save_to_file( data, file )
-        return file if File.exists?( file )
-
-        f = File.new( file, 'w' )
-        f.write( data )
-        f.close
-
+    def save_to_file( report, file )
+        report.save( file )
         Report.create(
-            :issue_count => get_issue_count( data ),
-            :host        => get_host( data ),
-            :filename    => File.basename( f.path, EXTENSION ),
-            :datestamp   => Time.now.asctime
+            :issue_count => get_issue_count( report ),
+            :host        => get_host( report ),
+            :filename    => File.basename( file, EXTENSION ),
+            :datestamp   => DateTime.now
         )
 
-        return f.path
+        return file
     end
 
     def convert( type, report )
-
         opts = {}
+
         classes[type].info[:options].each {
             |opt|
             opts[opt.name] = opt.default if opt.default
         }
+
         opts['outfile'] = get_tmp_outfile_name( type, report )
 
-        classes[type].new( unserialize( report ), opts ).run
+        classes[type].new( report, opts ).run
 
         content = File.read( opts['outfile'] )
         FileUtils.rm( opts['outfile'] )
         return content
     end
 
-
     def get_tmp_outfile_name( type, report )
-        tmpdir + get_filename( report ) + '.' + type
+        tmpdir + report_to_filename( report ) + '.' + type
     end
 
     def has_outfile?( options )
@@ -274,7 +269,6 @@ class ReportManager
         return false
     end
 
-
     def populate_available
         @@available ||= []
         return @@available if !@@available.empty?
@@ -291,7 +285,7 @@ class ReportManager
                     'rep_name'    => avail,
                     'description' => report_mgr[avail].info[:description],
                     'version'     => report_mgr[avail].info[:version],
-                    'author'      => report_mgr[avail].info[:author]
+                    'author'      => [report_mgr[avail].info[:author]].flatten
                 }
 
                 @@available_rep_classes[avail] = report_mgr[avail]