dstruct 0.0.1

data/lib/rex/io/datagram_abstraction.rb

@@ -0,0 +1,35 @@
+# -*- coding: binary -*-
+
+require 'socket'
+
+module Rex
+module IO
+
+###
+#
+# This class provides an abstraction to a datagram based
+# connection through the use of a datagram socketpair.
+#
+###
+module DatagramAbstraction
+
+  #
+  # Creates a streaming socket pair
+  #
+  def initialize_abstraction
+    self.lsock, self.rsock = Rex::Socket.udp_socket_pair()
+  end
+
+
+  # The left side of the stream (local)
+  attr_reader :lsock
+  # The right side of the stream (remote)
+  attr_reader :rsock
+
+protected
+  attr_writer :lsock
+  attr_writer :rsock
+
+end
+
+end; end

data/lib/rex/io/ring_buffer.rb

@@ -0,0 +1,369 @@
+# -*- coding: binary -*-
+#
+# This class implements a ring buffer with "cursors" in the form of sequence numbers.
+# To use this class, pass in a file descriptor and a ring size, the class will read
+# data from the file descriptor and store it in the ring. If the ring becomes full,
+# the oldest item will be overwritten. To emulate a stream interface, call read_data
+# to grab the last sequence number and any buffered data, call read_data again,
+# passing in the sequence number and all data newer than that sequence will be
+# returned, along with a new sequence to read from.
+#
+
+require 'rex/socket'
+
+module Rex
+module IO
+
+class RingBuffer
+
+  attr_accessor :queue     # The data queue, essentially an array of two-element arrays, containing a sequence and data buffer
+  attr_accessor :seq       # The next available sequence number
+  attr_accessor :fd        # The associated socket or IO object for this ring buffer
+  attr_accessor :size      # The number of available slots in the queue
+  attr_accessor :mutex     # The mutex locking access to the queue
+  attr_accessor :beg       # The index of the earliest data fragment in the ring
+  attr_accessor :cur       # The sequence number of the earliest data fragment in the ring
+  attr_accessor :monitor   # The thread handle of the built-in monitor when used
+  attr_accessor :monitor_thread_error  # :nodoc: #
+
+  #
+  # Create a new ring buffer
+  #
+  def initialize(socket, opts={})
+    self.size  = opts[:size] || (1024 * 4)
+    self.fd    = socket
+    self.seq   = 0
+    self.beg   = 0
+    self.cur   = 0
+    self.queue = Array.new( self.size )
+    self.mutex = Mutex.new
+  end
+
+  def inspect
+    "#<Rex::IO::RingBuffer @size=#{size} @fd=#{fd} @seq=#{seq} @beg=#{beg} @cur=#{cur}>"
+  end
+
+  #
+  # Start the built-in monitor, not called when used in a larger framework
+  #
+  def start_monitor
+    self.monitor = monitor_thread if not self.monitor
+  end
+
+  #
+  # Stop the built-in monitor
+  #
+  def stop_monitor
+    self.monitor.kill if self.monitor
+    self.monitor = nil
+  end
+
+  #
+  # The built-in monitor thread (normally unused with Metasploit)
+  #
+  def monitor_thread
+    Thread.new do
+      begin
+      while self.fd
+        buff = self.fd.get_once(-1, 1.0)
+        next if not buff
+        store_data(buff)
+      end
+      rescue ::Exception => e
+        self.monitor_thread_error = e
+      end
+    end
+  end
+
+  #
+  # Push data back into the associated stream socket. Logging must occur
+  # elsewhere, this function is simply a passthrough.
+  #
+  def put(data, opts={})
+    self.fd.put(data, opts={})
+  end
+
+  #
+  # The clear_data method wipes the ring buffer
+  #
+  def clear_data
+    self.mutex.synchronize do
+      self.seq   = 0
+      self.beg   = 0
+      self.cur   = 0
+      self.queue = Array.new( self.size )
+    end
+  end
+
+  #
+  # The store_data method is used to insert data into the ring buffer.
+  #
+  def store_data(data)
+    self.mutex.synchronize do
+      # self.cur points to the array index of queue containing the last item
+      # adding data will result in cur + 1 being used to store said data
+      # if cur is larger than size - 1, it will wrap back around. If cur
+      # is *smaller* beg, beg is increemnted to cur + 1 (and wrapped if
+      # necessary
+
+      loc = 0
+      if self.seq > 0
+        loc = ( self.cur + 1 ) % self.size
+
+        if loc <= self.beg
+          self.beg = (self.beg + 1) % self.size
+        end
+      end
+
+      self.queue[loc] = [self.seq += 1, data]
+      self.cur = loc
+    end
+  end
+
+  #
+  # The read_data method returns a two element array with the new reader cursor (a sequence number)
+  # and the returned data buffer (if any). A result of nil/nil indicates that no data is available
+  #
+  def read_data(ptr=nil)
+    self.mutex.synchronize do
+
+    # Verify that there is data in the queue
+    return [nil,nil] if not self.queue[self.beg]
+
+    # Configure the beginning read pointer (sequence number, not index)
+    ptr ||= self.queue[self.beg][0]
+    return [nil,nil] if not ptr
+
+    # If the pointer is below our baseline, we lost some data, so jump forward
+    if ptr < self.queue[self.beg][0]
+      ptr = self.queue[self.beg][0]
+    end
+
+    # Calculate how many blocks exist between the current sequence number
+    # and the requested pointer, this becomes the number of blocks we will
+    # need to read to satisfy the result. Due to the mutex block, we do
+    # not need to scan to find the sequence of the starting block or
+    # check the sequence of the ending block.
+    dis = self.seq - ptr
+
+    # If the requested sequnce number is less than our base pointer, it means
+    # that no new data is available and we should return empty.
+    return [nil,nil] if dis < 0
+
+    # Calculate the beginning block index and number of blocks to read
+    off = ptr - self.queue[self.beg][0]
+    set = (self.beg + off) % self.size
+
+
+    # Build the buffer by reading forward by the number of blocks needed
+    # and return the last read sequence number, plus one, as the new read
+    # pointer.
+    buff = ""
+    cnt  = 0
+    lst  = ptr
+    ptr.upto(self.seq) do |i|
+      block = self.queue[ (set + cnt) % self.size ]
+      lst,data = block[0],block[1]
+      buff += data
+      cnt += 1
+    end
+
+    return [lst + 1, buff]
+
+    end
+  end
+
+  #
+  # The base_sequence method returns the earliest sequence number in the queue. This is zero until
+  # all slots are filled and the ring rotates.
+  #
+  def base_sequence
+    self.mutex.synchronize do
+      return 0 if not self.queue[self.beg]
+      return self.queue[self.beg][0]
+    end
+  end
+
+  #
+  # The last_sequence method returns the "next" sequence number where new data will be
+  # available.
+  #
+  def last_sequence
+    self.seq
+  end
+
+  #
+  # The create_steam method assigns a IO::Socket compatible object to the ringer buffer
+  #
+  def create_stream
+    Stream.new(self)
+  end
+
+  #
+  # The select method returns when there is a chance of new data
+  # XXX: This is mostly useless and requires a rewrite to use a
+  #      real select or notify mechanism
+  #
+  def select
+    ::IO.select([ self.fd ], nil, [ self.fd ], 0.10)
+  end
+
+  #
+  # The wait method blocks until new data is available
+  #
+  def wait(seq)
+    nseq = nil
+    while not nseq
+      nseq,data = read_data(seq)
+      select
+    end
+  end
+
+  #
+  # The wait_for method blocks until new data is available or the timeout is reached
+  #
+  def wait_for(seq,timeout=1)
+    begin
+      ::Timeout.timeout(timeout) do
+        wait(seq)
+      end
+    rescue ::Timeout::Error
+    end
+  end
+
+  #
+  # This class provides a backwards compatible "stream" socket that uses
+  # the parents ring buffer.
+  #
+  class Stream
+    attr_accessor :ring
+    attr_accessor :seq
+    attr_accessor :buff
+
+    def initialize(ring)
+      self.ring = ring
+      self.seq  = ring.base_sequence
+      self.buff = ''
+    end
+
+    def read(len=nil)
+      if len and self.buff.length >= len
+        data = self.buff.slice!(0,len)
+        return data
+      end
+
+      while true
+        lseq, data = self.ring.read_data( self.seq )
+        return if not lseq
+
+        self.seq  = lseq
+        self.buff << data
+        if len
+          if self.buff.length >= len
+            return self.buff.slice!(0,len)
+          else
+            IO.select(nil, nil, nil, 0.25)
+            next
+          end
+        end
+
+        data = self.buff
+        self.buff = ''
+
+        return data
+
+        # Not reached
+        break
+      end
+
+    end
+
+    def write(data)
+      self.ring.write(data)
+    end
+  end
+
+end
+
+end
+end
+
+=begin
+
+server = Rex::Socket.create_tcp_server('LocalPort' => 0)
+lport  = server.getsockname[2]
+client = Rex::Socket.create_tcp('PeerHost' => '127.0.0.1', 'PeerPort' => lport)
+conn   = server.accept
+
+r = Rex::IO::RingBuffer.new(conn, {:size => 1024*1024})
+client.put("1")
+client.put("2")
+client.put("3")
+
+s,d = r.read_data
+
+client.put("4")
+client.put("5")
+client.put("6")
+s,d = r.read_data(s)
+
+client.put("7")
+client.put("8")
+client.put("9")
+s,d = r.read_data(s)
+
+client.put("0")
+s,d = r.read_data(s)
+
+test_counter = 11
+1.upto(100) do
+  client.put( "X" )
+  test_counter += 1
+end
+
+sleep(1)
+
+s,d = r.read_data
+p s
+p d
+
+fdata = ''
+File.open("/bin/ls", "rb") do |fd|
+  fdata = fd.read(fd.stat.size)
+  fdata = fdata * 10
+  client.put(fdata)
+end
+
+sleep(1)
+
+s,vdata = r.read_data(s)
+
+if vdata != fdata
+  puts "DATA FAILED"
+else
+  puts "DATA VERIFIED"
+end
+
+r.clear_data
+
+a = r.create_stream
+b = r.create_stream
+
+client.put("ABC123")
+sleep(1)
+
+p a.read
+p b.read
+
+client.put("$$$$$$")
+sleep(1)
+
+p a.read
+p b.read
+
+c = r.create_stream
+p c.read
+
+=end
+
+

data/lib/rex/io/stream.rb

@@ -0,0 +1,312 @@
+# -*- coding: binary -*-
+require 'rex/sync/thread_safe'
+
+module Rex
+module IO
+
+###
+#
+# This mixin is an abstract representation of a streaming connection.  Streams
+# extend classes that must implement the following methods:
+#
+#   syswrite(buffer)
+#   sysread(length)
+#   shutdown(how)
+#   close
+#   peerinfo
+#   localinfo
+#
+###
+module Stream
+
+  ##
+  #
+  # Abstract methods
+  #
+  ##
+
+  #
+  # This method writes the supplied buffer to the stream.  This method
+  # intelligent reduces the size of supplied buffers so that ruby doesn't get
+  # into a potential global thread blocking state when used on blocking
+  # sockets.  That is, this method will send the supplied buffer in chunks
+  # of, at most, 32768 bytes.
+  #
+  def write(buf, opts = {})
+    total_sent   = 0
+    total_length = buf.length
+    block_size   = 32768
+
+    begin
+      while( total_sent < total_length )
+        s = Rex::ThreadSafe.select( nil, [ fd ], nil, 0.2 )
+        if( s == nil || s[0] == nil )
+          next
+        end
+        data = buf[total_sent, block_size]
+        sent = fd.write_nonblock( data )
+        if sent > 0
+          total_sent += sent
+        end
+      end
+    rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
+      # Sleep for a half a second, or until we can write again
+      Rex::ThreadSafe.select( nil, [ fd ], nil, 0.5 )
+      # Decrement the block size to handle full sendQs better
+      block_size = 1024
+      # Try to write the data again
+      retry
+    rescue ::IOError, ::Errno::EPIPE
+      return nil
+    end
+
+    total_sent
+  end
+
+  #
+  # This method reads data of the supplied length from the stream.
+  #
+  def read(length = nil, opts = {})
+
+    begin
+      return fd.read_nonblock( length )
+    rescue ::Errno::EAGAIN, ::Errno::EWOULDBLOCK
+      # Sleep for a half a second, or until we can read again
+      Rex::ThreadSafe.select( [ fd ], nil, nil, 0.5 )
+      # Decrement the block size to handle full sendQs better
+      retry
+    rescue ::IOError, ::Errno::EPIPE
+      return nil
+    end
+  end
+
+  #
+  # Polls the stream to see if there is any read data available.  Returns
+  # true if data is available for reading, otherwise false is returned.
+  #
+  def has_read_data?(timeout = nil)
+
+    # Allow a timeout of "0" that waits almost indefinitely for input, this
+    # mimics the behavior of Rex::ThreadSafe.select() and fixes some corner
+    # cases of unintentional no-wait timeouts.
+    timeout = 3600 if (timeout and timeout == 0)
+
+    begin
+      if ((rv = ::IO.select([ fd ], nil, nil, timeout)) and
+          (rv[0]) and
+          (rv[0][0] == fd))
+        true
+      else
+        false
+      end
+    rescue ::Errno::EBADF, ::Errno::ENOTSOCK
+      raise ::EOFError
+    rescue StreamClosedError, ::IOError, ::EOFError, ::Errno::EPIPE
+      #  Return false if the socket is dead
+      return false
+    end
+  end
+
+  #
+  # This method returns the selectable file descriptor, or self by default.
+  #
+  def fd
+    self
+  end
+
+  ##
+  #
+  # Common methods
+  #
+  ##
+
+  #
+  # This method writes the supplied buffer to the stream by calling the write
+  # routine.
+  #
+  def <<(buf)
+    return write(buf.to_s)
+  end
+
+  #
+  # This method calls get_once() to read pending data from the socket
+  #
+  def >>
+    get_once
+  end
+
+  #
+  # This method writes to the stream, optionally timing out after a period of
+  # time.
+  #
+  def timed_write(buf, wait = def_write_timeout, opts = {})
+    if (wait and wait > 0)
+      Timeout.timeout(wait) {
+        return write(buf, opts)
+      }
+    else
+      return write(buf, opts)
+    end
+  end
+
+  #
+  # This method reads from the stream, optionally timing out after a period
+  # of time.
+  #
+  def timed_read(length = nil, wait = def_read_timeout, opts = {})
+    if (wait and wait > 0)
+      Timeout.timeout(wait) {
+        return read(length, opts)
+      }
+    else
+      return read(length, opts)
+    end
+  end
+
+  #
+  # This method writes the full contents of the supplied buffer, optionally
+  # with a timeout.
+  #
+  def put(buf, opts = {})
+    return 0 if (buf == nil or buf.length == 0)
+
+    send_len = buf.length
+    send_idx = 0
+    wait     = opts['Timeout'] || 0
+
+    # Keep writing until our send length drops to zero
+    while (send_idx < send_len)
+      curr_len  = timed_write(buf[send_idx, buf.length-send_idx], wait, opts)
+
+      # If the write operation failed due to an IOError, then we fail.
+      return buf.length - send_len if (curr_len == nil)
+
+      send_len -= curr_len
+      send_idx += curr_len
+    end
+
+    return buf.length - send_len
+  end
+
+
+  #
+  # This method emulates the behavior of Pex::Socket::Recv in MSF2
+  #
+  def get_once(length = -1, timeout = def_read_timeout)
+
+    if (has_read_data?(timeout) == false)
+      return nil
+    end
+
+    bsize = (length == -1) ? def_block_size : length
+    data  = read(bsize)
+    raise EOFError if data.nil?
+    data
+  end
+
+  #
+  # This method reads as much data as it can from the wire given a maximum
+  # timeout.
+  #
+  def get(timeout = nil, ltimeout = def_read_loop_timeout, opts = {})
+    # For those people who are used to being able to use a negative timeout!
+    if (timeout and timeout.to_i < 0)
+      timeout = nil
+    end
+
+    # No data in the first place? bust.
+    if (has_read_data?(timeout) == false)
+      return nil
+    end
+
+    buf = ""
+    lps = 0
+    eof = false
+
+    # Keep looping until there is no more data to be gotten..
+    while (has_read_data?(ltimeout) == true)
+      # Catch EOF errors so that we can handle them properly.
+      begin
+        temp = read(def_block_size)
+      rescue EOFError
+        eof = true
+      end
+
+      # If we read zero bytes and we had data, then we've hit EOF
+      if (temp and temp.length == 0)
+        eof = true
+      end
+
+      # If we reached EOF and there are no bytes in the buffer we've been
+      # reading into, then throw an EOF error.
+      if (eof)
+        # If we've already read at least some data, then it's time to
+        # break out and let it be processed before throwing an EOFError.
+        if (buf.length > 0)
+          break
+        else
+          raise EOFError
+        end
+      end
+
+      break if (temp == nil or temp.empty? == true)
+
+      buf += temp
+      lps += 1
+
+      break if (lps >= def_max_loops)
+    end
+
+    # Return the entire buffer we read in
+    return buf
+  end
+
+  ##
+  #
+  # Defaults
+  #
+  ##
+
+  #
+  # The default number of seconds to wait for a write operation to timeout.
+  #
+  def def_write_timeout
+    10
+  end
+
+  #
+  # The default number of seconds to wait for a read operation to timeout.
+  #
+  def def_read_timeout
+    10
+  end
+
+  #
+  # The default number of seconds to wait while in a read loop after read
+  # data has been found.
+  #
+  def def_read_loop_timeout
+    0.1
+  end
+
+  #
+  # The maximum number of read loops to perform before returning to the
+  # caller.
+  #
+  def def_max_loops
+    1024
+  end
+
+  #
+  # The default block size to read in chunks from the wire.
+  #
+  def def_block_size
+    16384
+  end
+
+protected
+
+end
+
+end end
+