require 'zlib'
require 'net/ssh/transport/cipher_factory'
require 'net/ssh/transport/hmac'

module Net; module SSH; module Transport

  # Encapsulates state information about one end of an SSH connection. Such
  # state includes the packet sequence number, the algorithms in use, how
  # many packets and blocks have been processed since the last reset, and so
  # forth. This class will never be instantiated directly, but is used as
  # part of the internal state of the PacketStream module.
  class State
    # The socket object that owns this state object.
    attr_reader :socket

    # The next packet sequence number for this socket endpoint.
    attr_reader :sequence_number

    # The hmac algorithm in use for this endpoint.
    attr_reader :hmac

    # The compression algorithm in use for this endpoint.
    attr_reader :compression

    # The compression level to use when compressing data (or nil, for the default).
    attr_reader :compression_level

    # The number of packets processed since the last call to #reset!
    attr_reader :packets

    # The number of data blocks processed since the last call to #reset!
    attr_reader :blocks

    # The cipher algorithm in use for this socket endpoint.
    attr_reader :cipher

    # The block size for the cipher
    attr_reader :block_size

    # The role that this state plays (either :client or :server)
    attr_reader :role

    # The maximum number of packets that this endpoint wants to process before
    # needing a rekey.
    attr_accessor :max_packets

    # The maximum number of blocks that this endpoint wants to process before
    # needing a rekey.
    attr_accessor :max_blocks

    # The user-specified maximum number of bytes that this endpoint ought to
    # process before needing a rekey.
    attr_accessor :rekey_limit

    # Creates a new state object, belonging to the given socket. Initializes
    # the algorithms to "none".
    def initialize(socket, role)
      @socket = socket
      @role = role
      @sequence_number = @packets = @blocks = 0
      @cipher = CipherFactory.get("none")
      @block_size = 8
      @hmac = HMAC.get("none")
      @compression = nil
      @compressor = @decompressor = nil
      @next_iv = ""
    end

    # A convenience method for quickly setting multiple values in a single
    # command.
    def set(values)
      values.each do |key, value|
        instance_variable_set("@#{key}", value)
      end
      reset!
    end

    def update_cipher(data)
      result = cipher.update(data)
      update_next_iv(role == :client ? result : data)
      return result
    end

    def final_cipher
      result = cipher.final
      update_next_iv(role == :client ? result : "", true)
      return result
    end

    # Increments the counters. The sequence number is incremented (and remapped
    # so it always fits in a 32-bit integer). The number of packets and blocks
    # are also incremented.
    def increment(packet_length)
      @sequence_number = (@sequence_number + 1) & 0xFFFFFFFF
      @packets += 1
      @blocks += (packet_length + 4) / @block_size
    end

    # The compressor object to use when compressing data. This takes into account
    # the desired compression level.
    def compressor
      @compressor ||= Zlib::Deflate.new(compression_level || Zlib::DEFAULT_COMPRESSION)
    end

    # The decompressor object to use when decompressing data.
    def decompressor
      @decompressor ||= Zlib::Inflate.new(nil)
    end

    # Returns true if data compression/decompression is enabled. This will
    # return true if :standard compression is selected, or if :delayed
    # compression is selected and the :authenticated hint has been received
    # by the socket.
    def compression?
      compression == :standard || (compression == :delayed && socket.hints[:authenticated])
    end

    # Compresses the data. If no compression is in effect, this will just return
    # the data unmodified, otherwise it uses #compressor to compress the data.
    def compress(data)
      data = data.to_s
      return data unless compression?
      compressor.deflate(data, Zlib::SYNC_FLUSH)
    end

    # Deompresses the data. If no compression is in effect, this will just return
    # the data unmodified, otherwise it uses #decompressor to decompress the data.
    def decompress(data)
      data = data.to_s
      return data unless compression?
      decompressor.inflate(data)
    end

    # Resets the counters on the state object, but leaves the sequence_number
    # unchanged. It also sets defaults for and recomputes the max_packets and
    # max_blocks values.
    def reset!
      @packets = @blocks = 0

      @max_packets ||= 1 << 31

      @block_size = cipher.name == "RC4" ? 8 : cipher.block_size

      if max_blocks.nil?
        # cargo-culted from openssh. the idea is that "the 2^(blocksize*2)
        # limit is too expensive for 3DES, blowfish, etc., so enforce a 1GB
        # limit for small blocksizes."
        if @block_size >= 16
          @max_blocks = 1 << (@block_size * 2)
        else
          @max_blocks = (1 << 30) / @block_size
        end

        # if a limit on the # of bytes has been given, convert that into a
        # minimum number of blocks processed.

        if rekey_limit
          @max_blocks = [@max_blocks, rekey_limit / @block_size].min
        end
      end

      cleanup
    end

    # Closes any the compressor and/or decompressor objects that have been
    # instantiated.
    def cleanup
      if @compressor
        @compressor.finish if !@compressor.finished?
        @compressor.close
      end

      if @decompressor
        # we call reset here so that we don't get warnings when we try to
        # close the decompressor
        @decompressor.reset
        @decompressor.close
      end

      @compressor = @decompressor = nil
    end

    # Returns true if the number of packets processed exceeds the maximum
    # number of packets, or if the number of blocks processed exceeds the
    # maximum number of blocks.
    def needs_rekey?
      max_packets && packets > max_packets ||
      max_blocks && blocks > max_blocks
    end

    private

      def update_next_iv(data, reset=false)
        @next_iv << data
        @next_iv = @next_iv[-cipher.iv_len..-1]

        if reset
          cipher.reset
          cipher.iv = @next_iv
        end

        return data
      end
  end

end; end; end