# -*- 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