require "socket"
require "delegate"
require "timeout"
require "English"
require "net/telnet.rb"

module Net

  class Telnet_Patched < Telnet

    TERM_VT100 = "vt100"

    # Preprocess received data from the host.
    #
    # Performs newline conversion and detects telnet command sequences.
    # Called automatically by #waitfor().  You should only use this
    # method yourself if you have read input directly using sysread()
    # or similar, and even then only if in telnet mode.
    def preprocess(string)
      # combine CR+NULL into CR
      string = string.gsub(/#{CR}#{NULL}/no, CR) if @options["Telnetmode"]

      # combine EOL into "\n"
      string = string.gsub(/#{EOL}/no, "\n") unless @options["Binmode"]

      string.gsub(/#{IAC}(
                   [#{IAC}#{AO}#{AYT}#{DM}#{IP}#{NOP}]|
                   [#{DO}#{DONT}#{WILL}#{WONT}]
                     [#{OPT_BINARY}-#{OPT_NEW_ENVIRON}#{OPT_EXOPL}]|
                   #{SB}[^#{IAC}]*#{IAC}#{SE}
                 )/xno) do
        if    IAC == $1  # handle escaped IAC characters
          IAC
        elsif AYT == $1  # respond to "IAC AYT" (are you there)
          self.write("nobody here but us pigeons" + EOL)
          ''
        elsif DO[0] == $1[0]  # respond to "IAC DO x"
          if OPT_BINARY[0] == $1[1]
            @telnet_option["BINARY"] = true
            self.write(IAC + WILL + OPT_BINARY)
          elsif OPT_TTYPE[0] == $1[1]
            # NO! I want to OPT_TTYPE!
            self.write(IAC + WILL + OPT_TTYPE)
          else
            self.write(IAC + WONT + $1[1..1])
            #self.write(IAC + WILL + $1[1..1])
          end
          ''
        elsif DONT[0] == $1[0]  # respond to "IAC DON'T x" with "IAC WON'T x"
          self.write(IAC + WONT + $1[1..1])
          ''
        elsif WILL[0] == $1[0]  # respond to "IAC WILL x"
          if    OPT_BINARY[0] == $1[1]
            self.write(IAC + DO + OPT_BINARY)
          elsif OPT_ECHO[0] == $1[1]
            self.write(IAC + DO + OPT_ECHO)
          elsif OPT_SGA[0]  == $1[1]
            @telnet_option["SGA"] = true
            self.write(IAC + DO + OPT_SGA)
          else
            self.write(IAC + DONT + $1[1..1])
          end
          ''
        elsif WONT[0] == $1[0]  # respond to "IAC WON'T x"
          if    OPT_ECHO[0] == $1[1]
            self.write(IAC + DONT + OPT_ECHO)
          elsif OPT_SGA[0]  == $1[1]
            @telnet_option["SGA"] = false
            self.write(IAC + DONT + OPT_SGA)
          else
            self.write(IAC + DONT + $1[1..1])
          end
          ''
        elsif SB[0] == $1[0]
          if OPT_TTYPE[0] == $1[1]
            # want to send 0xfffa 0x1800 <vt100 ascii> 0xfff0
            self.write(IAC + SB + OPT_TTYPE + OPT_BINARY + TERM_VT100 + IAC + SE)
          end
        else
          ''
        end
      end
    end # preprocess

    # Read data from the host until a certain sequence is matched.
    #
    # If a block is given, the received data will be yielded as it
    # is read in (not necessarily all in one go), or nil if EOF
    # occurs before any data is received.  Whether a block is given
    # or not, all data read will be returned in a single string, or again
    # nil if EOF occurs before any data is received.  Note that
    # received data includes the matched sequence we were looking for.
    #
    # +options+ can be either a regular expression or a hash of options.
    # If a regular expression, this specifies the data to wait for.
    # If a hash, this can specify the following options:
    #
    # Match:: a regular expression, specifying the data to wait for.
    # Prompt:: as for Match; used only if Match is not specified.
    # String:: as for Match, except a string that will be converted
    #          into a regular expression.  Used only if Match and
    #          Prompt are not specified.
    # Timeout:: the number of seconds to wait for data from the host
    #           before raising a TimeoutError.  If set to false,
    #           no timeout will occur.  If not specified, the
    #           Timeout option value specified when this instance
    #           was created will be used, or, failing that, the
    #           default value of 10 seconds.
    # Waittime:: the number of seconds to wait after matching against
    #            the input data to see if more data arrives.  If more
    #            data arrives within this time, we will judge ourselves
    #            not to have matched successfully, and will continue
    #            trying to match.  If not specified, the Waittime option
    #            value specified when this instance was created will be
    #            used, or, failing that, the default value of 0 seconds,
    #            which means not to wait for more input.

    def waitfor(options) # :yield: recvdata
      time_out = @options["Timeout"]
      waittime = @options["Waittime"]

      if options.kind_of?(Hash)
        prompt   = if options.has_key?("Match")
                     options["Match"]
                   elsif options.has_key?("Prompt")
                     options["Prompt"]
                   elsif options.has_key?("String")
                     Regexp.new( Regexp.quote(options["String"]) )
                   end
        time_out = options["Timeout"]  if options.has_key?("Timeout")
        waittime = options["Waittime"] if options.has_key?("Waittime")
      else
        prompt = options
      end

      if time_out == false
        time_out = nil
      end

      line = ''
      buf = ''
      rest = ''
      until(prompt === line and not IO::select([@sock], nil, nil, waittime))
        unless IO::select([@sock], nil, nil, time_out)
          raise TimeoutError, "timed out while waiting for more data"
        end
        begin
          c = @sock.readpartial(1024 * 1024)
          @dumplog.log_dump('<', c) if @options.has_key?("Dump_log")
          if @options["Telnetmode"]
            c = rest + c
            if Integer(c.rindex(/#{IAC}#{SE}/no)) <
               Integer(c.rindex(/#{IAC}#{SB}/no))
              buf = preprocess(c[0 ... c.rindex(/#{IAC}#{SB}/no)])
              rest = c[c.rindex(/#{IAC}#{SB}/no) .. -1]
            elsif pt = c.rindex(/#{IAC}[^#{IAC}#{AO}#{AYT}#{DM}#{IP}#{NOP}]?\z/no)
              # without really knowing why, if we chomp "rest" off the end of
              # the string, we lose the {IAC SE} at the end of the suboption
              # which means the tricky regexp in preprocess for handling
              # {IAC SB <stuff> IAC SE} will fail.

              #buf = preprocess(c[0 ... pt])
              #rest = c[pt .. -1]
              buf = preprocess(c)       #just process it, fool
            else
              buf = preprocess(c)
              rest = ''
            end
         else
           # Not Telnetmode.
           #
           # We cannot use preprocess() on this data, because that
           # method makes some Telnetmode-specific assumptions.
           buf = c
           buf.gsub!(/#{EOL}/no, "\n") unless @options["Binmode"]
           rest = ''
          end
          @log.print(buf) if @options.has_key?("Output_log")
          line += buf
          yield buf if block_given?
        rescue EOFError # End of file reached
          if line == ''
            line = nil
            yield nil if block_given?
          end
          break
        end
      end
      line
    end
  end # class Telnet_Patched
end # module Net