dstruct 0.0.1

data/lib/rex/io/stream_abstraction.rb

@@ -0,0 +1,209 @@
+# -*- coding: binary -*-
+
+require 'socket'
+require 'fcntl'
+
+module Rex
+module IO
+
+###
+#
+# This class provides an abstraction to a stream based
+# connection through the use of a streaming socketpair.
+#
+###
+module StreamAbstraction
+
+  ###
+  #
+  # Extension information for required Stream interface.
+  #
+  ###
+  module Ext
+
+    #
+    # Initializes peer information.
+    #
+    def initinfo(peer,local)
+      @peer = peer
+      @local = local
+    end
+
+    #
+    # Symbolic peer information.
+    #
+    def peerinfo
+      (@peer || "Remote Pipe")
+    end
+
+    #
+    # Symbolic local information.
+    #
+    def localinfo
+      (@local || "Local Pipe")
+    end
+  end
+
+  #
+  # This method creates a streaming socket pair and initializes it.
+  #
+  def initialize_abstraction
+    self.lsock, self.rsock = Rex::Socket.tcp_socket_pair()
+    self.lsock.extend(Rex::IO::Stream)
+    self.lsock.extend(Ext)
+    self.rsock.extend(Rex::IO::Stream)
+
+    self.monitor_rsock
+  end
+
+  #
+  # This method cleans up the abstraction layer.
+  #
+  def cleanup_abstraction
+    self.lsock.close if (self.lsock)
+    self.rsock.close if (self.rsock)
+
+    self.lsock = nil
+    self.rsock = nil
+  end
+
+  #
+  # Low-level write to the local side.
+  #
+  def syswrite(buffer)
+    lsock.syswrite(buffer)
+  end
+
+  #
+  # Low-level read from the local side.
+  #
+  def sysread(length)
+    lsock.sysread(length)
+  end
+
+  #
+  # Shuts down the local side of the stream abstraction.
+  #
+  def shutdown(how)
+    lsock.shutdown(how)
+  end
+
+  #
+  # Closes both sides of the stream abstraction.
+  #
+  def close
+    cleanup_abstraction
+  end
+
+  #
+  # Symbolic peer information.
+  #
+  def peerinfo
+    "Remote-side of Pipe"
+  end
+
+  #
+  # Symbolic local information.
+  #
+  def localinfo
+    "Local-side of Pipe"
+  end
+
+  #
+  # The left side of the stream.
+  #
+  attr_reader :lsock
+  #
+  # The right side of the stream.
+  #
+  attr_reader :rsock
+
+protected
+
+  def monitor_rsock
+    self.monitor_thread = Rex::ThreadFactory.spawn("StreamMonitorRemote", false) {
+      loop do
+        closed = false
+        buf    = nil
+
+        if not self.rsock
+          wlog("monitor_rsock: the remote socket is nil, exiting loop")
+          break
+        end
+
+        begin
+          s = Rex::ThreadSafe.select( [ self.rsock ], nil, nil, 0.2 )
+          if( s == nil || s[0] == nil )
+            next
+          end
+        rescue Exception => e
+          wlog("monitor_rsock: exception during select: #{e.class} #{e}")
+          closed = true
+        end
+
+        if( closed == false )
+          begin
+            buf = self.rsock.sysread( 32768 )
+            if buf == nil
+              closed = true
+              wlog("monitor_rsock: closed remote socket due to nil read")
+            end
+          rescue EOFError => e
+            closed = true
+            dlog("monitor_rsock: EOF in rsock")
+          rescue ::Exception => e
+            closed = true
+            wlog("monitor_rsock: exception during read: #{e.class} #{e}")
+          end
+        end
+
+        if( closed == false )
+          total_sent   = 0
+          total_length = buf.length
+          while( total_sent < total_length )
+            begin
+              data = buf[total_sent, buf.length]
+
+              # Note that this must be write() NOT syswrite() or put() or anything like it.
+              # Using syswrite() breaks SSL streams.
+              sent = self.write( data )
+
+              # sf: Only remove the data off the queue is write was successfull.
+              #     This way we naturally perform a resend if a failure occured.
+              #     Catches an edge case with meterpreter TCP channels where remote send
+              #     failes gracefully and a resend is required.
+              if (sent.nil?)
+                closed = true
+                wlog("monitor_rsock: failed writing, socket must be dead")
+                break
+              elsif (sent > 0)
+                total_sent += sent
+              end
+            rescue ::IOError, ::EOFError => e
+              closed = true
+              wlog("monitor_rsock: exception during write: #{e.class} #{e}")
+              break
+            end
+          end
+        end
+
+        if( closed )
+          begin
+            self.close_write if self.respond_to?('close_write')
+          rescue IOError
+          end
+          break
+        end
+      end
+    }
+  end
+
+protected
+  attr_accessor :monitor_thread
+  attr_writer :lsock
+  attr_writer :rsock
+
+end
+
+end; end
+

data/lib/rex/io/stream_server.rb

@@ -0,0 +1,221 @@
+# -*- coding: binary -*-
+require 'thread'
+
+module Rex
+module IO
+
+###
+#
+# This mixin provides the framework and interface for implementing a streaming
+# server that can listen for and accept stream client connections.  Stream
+# servers extend this class and are required to implement the following
+# methods:
+#
+#   accept
+#   fd
+#
+###
+module StreamServer
+
+  ##
+  #
+  # Abstract methods
+  #
+  ##
+
+  ##
+  #
+  # Default server monitoring and client management implementation follows
+  # below.
+  #
+  ##
+
+  #
+  # This callback is notified when a client connects.
+  #
+  def on_client_connect(client)
+    if (on_client_connect_proc)
+      on_client_connect_proc.call(client)
+    end
+  end
+
+  #
+  # This callback is notified when a client connection has data that needs to
+  # be processed.
+  #
+  def on_client_data(client)
+    if (on_client_data_proc)
+      on_client_data_proc.call(client)
+    end
+  end
+
+  #
+  # This callback is notified when a client connection has closed.
+  #
+  def on_client_close(client)
+    if (on_client_close_proc)
+      on_client_close_proc.call(client)
+    end
+  end
+
+  #
+  # Start monitoring the listener socket for connections and keep track of
+  # all client connections.
+  #
+  def start
+    self.clients = []
+    self.client_waiter = ::Queue.new
+
+    self.listener_thread = Rex::ThreadFactory.spawn("StreamServerListener", false) {
+      monitor_listener
+    }
+    self.clients_thread = Rex::ThreadFactory.spawn("StreamServerClientMonitor", false) {
+      monitor_clients
+    }
+  end
+
+  #
+  # Terminates the listener monitoring threads and closes all active clients.
+  #
+  def stop
+    self.listener_thread.kill
+    self.clients_thread.kill
+
+    self.clients.each { |cli|
+      close_client(cli)
+    }
+  end
+
+  #
+  # This method closes a client connection and cleans up the resources
+  # associated with it.
+  #
+  def close_client(client)
+    if (client)
+      clients.delete(client)
+
+      begin
+        client.close
+      rescue IOError
+      end
+    end
+  end
+
+  #
+  # This method waits on the server listener thread
+  #
+  def wait
+    self.listener_thread.join if self.listener_thread
+  end
+
+  ##
+  #
+  # Callback procedures.
+  #
+  ##
+
+  #
+  # This callback procedure can be set and will be called when new clients
+  # connect.
+  #
+  attr_accessor :on_client_connect_proc
+  #
+  # This callback procedure can be set and will be called when clients
+  # have data to be processed.
+  #
+  attr_accessor :on_client_data_proc
+  #
+  # This callback procedure can be set and will be called when a client
+  # disconnects from the server.
+  #
+  attr_accessor :on_client_close_proc
+
+  attr_accessor :clients # :nodoc:
+  attr_accessor :listener_thread, :clients_thread # :nodoc:
+  attr_accessor :client_waiter
+
+protected
+
+  #
+  # This method monitors the listener socket for new connections and calls
+  # the +on_client_connect+ callback routine.
+  #
+  def monitor_listener
+
+    while true
+      begin
+        cli = accept
+        if not cli
+          elog("The accept() returned nil in stream server listener monitor:  #{fd.inspect}")
+          ::IO.select(nil, nil, nil, 0.10)
+          next
+        end
+
+        # Append to the list of clients
+        self.clients << cli
+
+        # Initialize the connection processing
+        on_client_connect(cli)
+
+        # Notify the client monitor
+        self.client_waiter.push(cli)
+
+      # Skip exceptions caused by accept() [ SSL ]
+      rescue ::EOFError, ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED
+      rescue ::Interrupt
+        raise $!
+      rescue ::Exception
+        elog("Error in stream server server monitor: #{$!}")
+        rlog(ExceptionCallStack)
+        break
+      end
+    end
+  end
+
+  #
+  # This method monitors client connections for data and calls the
+  # +on_client_data+ routine when new data arrives.
+  #
+  def monitor_clients
+    begin
+
+      # Wait for a notify if our client list is empty
+      if (clients.length == 0)
+        self.client_waiter.pop
+        next
+      end
+
+      sd = Rex::ThreadSafe.select(clients, nil, nil, nil)
+
+      sd[0].each { |cfd|
+        begin
+          on_client_data(cfd)
+        rescue ::EOFError, ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED
+          on_client_close(cfd)
+          close_client(cfd)
+        rescue ::Interrupt
+          raise $!
+        rescue ::Exception
+          close_client(cfd)
+          elog("Error in stream server client monitor: #{$!}")
+          rlog(ExceptionCallStack)
+
+        end
+      }
+
+    rescue ::Rex::StreamClosedError => e
+      # Remove the closed stream from the list
+      clients.delete(e.stream)
+    rescue ::Interrupt
+      raise $!
+    rescue ::Exception
+      elog("Error in stream server client monitor: #{$!}")
+      rlog(ExceptionCallStack)
+    end while true
+  end
+
+end
+
+end
+end
+

data/lib/rex/job_container.rb

@@ -0,0 +1,200 @@
+# -*- coding: binary -*-
+module Rex
+
+###
+#
+# This class is the concrete representation of an abstract job.
+#
+###
+class Job
+
+  #
+  # Creates an individual job instance and initializes it with the supplied
+  # parameters.
+  #
+  def initialize(container, jid, name, ctx, run_proc, clean_proc)
+    self.container  = container
+    self.jid        = jid
+    self.name       = name
+    self.run_proc   = run_proc
+    self.clean_proc = clean_proc
+    self.ctx        = ctx
+    self.start_time = nil
+  end
+
+  #
+  # Runs the job in the context of its own thread if the async flag is false.
+  # Otherwise, the job is run inline.
+  #
+  def start(async = false)
+    self.start_time = Time.now
+    if (async)
+      self.job_thread = Rex::ThreadFactory.spawn("JobID(#{jid})-#{name}", false) {
+        # Deschedule our thread momentarily
+        ::IO.select(nil, nil, nil, 0.01)
+
+        begin
+          run_proc.call(ctx)
+        ensure
+          clean_proc.call(ctx)
+          container.remove_job(self)
+        end
+      }
+    else
+      begin
+        run_proc.call(ctx)
+      rescue ::Exception
+        container.stop_job(jid)
+        raise $!
+      end
+    end
+  end
+
+  #
+  # Stops the job if it's currently running and calls its cleanup procedure
+  #
+  def stop
+    if (self.job_thread)
+      self.job_thread.kill
+      self.job_thread = nil
+    end
+
+    clean_proc.call(ctx) if (clean_proc)
+  end
+
+  #
+  # The name of the job.
+  #
+  attr_reader :name
+
+  #
+  # The job identifier as assigned by the job container.
+  #
+  attr_reader :jid
+
+  #
+  # The time at which this job was started.
+  #
+  attr_reader :start_time
+
+  #
+  # Some job context.
+  #
+  attr_reader :ctx
+
+protected
+
+  attr_writer   :name #:nodoc:
+  attr_writer   :jid #:nodoc:
+  attr_accessor :job_thread #:nodoc:
+  attr_accessor :container #:nodoc:
+  attr_accessor :run_proc #:nodoc:
+  attr_accessor :clean_proc #:nodoc:
+  attr_writer   :ctx #:nodoc:
+  attr_writer   :start_time #:nodoc:
+
+end
+
+###
+#
+# This class contains zero or more abstract jobs that can be enumerated and
+# stopped in a generic fashion.  This is used to provide a mechanism for
+# keeping track of arbitrary contexts that may or may not require a dedicated
+# thread.
+#
+###
+class JobContainer < Hash
+
+  def initialize
+    self.job_id_pool = 0
+  end
+
+  #
+  # Adds an already running task as a symbolic job to the container.
+  #
+  def add_job(name, ctx, run_proc, clean_proc)
+    real_name = name
+    count     = 0
+    jid       = job_id_pool
+
+    self.job_id_pool += 1
+
+    # If we were not supplied with a job name, pick one from the hat
+    if (real_name == nil)
+      real_name = '#' + jid.to_s
+    end
+
+    # Find a unique job name
+    while (j = self[real_name])
+      real_name  = name + " #{count}"
+      count     += 1
+    end
+
+    j = Job.new(self, jid, real_name, ctx, run_proc, clean_proc)
+
+    self[jid.to_s] = j
+  end
+
+  #
+  # Starts a job using the supplied name and run/clean procedures.
+  #
+  def start_job(name, ctx, run_proc, clean_proc = nil)
+    j = add_job(name, ctx, run_proc, clean_proc)
+    j.start
+
+    j.jid
+  end
+
+  #
+  # Starts a background job that doesn't call the cleanup routine or run
+  # the run_proc in its own thread.  Rather, the run_proc is called
+  # immediately and the clean_proc is never called until the job is removed
+  # from the job container.
+  #
+  def start_bg_job(name, ctx, run_proc, clean_proc = nil, async = true)
+    j = add_job(name, ctx, run_proc, clean_proc)
+    j.start(async)
+
+    j.jid
+  end
+
+  #
+  # Stops the job with the supplied name and forces it to cleanup.  Stopping
+  # the job also leads to its removal.
+  #
+  def stop_job(jid)
+    if (j = self[jid.to_s])
+      j.stop
+
+      remove_job(j)
+    end
+  end
+
+  #
+  # Removes a job that was previously running.  This is typically called when
+  # a job completes its task.
+  #
+  def remove_job(inst)
+    self.delete(inst.jid.to_s)
+  end
+
+  #
+  # Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
+  #    "can't add a new key into hash during iteration"
+  #
+  def each(&block)
+    list = []
+    self.keys.sort.each do |sidx|
+      list << [sidx, self[sidx]]
+    end
+    list.each(&block)
+  end
+
+protected
+
+  attr_accessor :job_id_pool # :nodoc:
+
+end
+
+end
+