# -*- coding: binary -*-
require 'msf/core/exploit/tcp_server'

module Msf

module EvasiveTCP
  attr_accessor :_send_size, :_send_delay, :evasive

  def denagle
    begin
      setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
    rescue ::Exception
    end
  end

  def write(buf, opts={})

    return super(buf, opts) if not @evasive

    ret = 0
    idx = 0
    len = @_send_size || buf.length

    while(idx < buf.length)

      if(@_send_delay and idx > 0)
        ::IO.select(nil, nil, nil, @_send_delay)
      end

      pkt = buf[idx, len]

      res = super(pkt, opts)
      flush()

      idx += len
      ret += res if res
    end
    ret
  end
end

###
#
# This module provides methods for establish a connection to a remote host and
# communicating with it.
#
###
module Exploit::Remote::Tcp

  #
  # Initializes an instance of an exploit module that exploits a
  # vulnerability in a TCP server.
  #
  def initialize(info = {})
    super

    register_options(
      [
        Opt::RHOST,
        Opt::RPORT
      ], Msf::Exploit::Remote::Tcp
    )

    register_advanced_options(
      [
        OptBool.new('SSL',        [ false, 'Negotiate SSL/TLS for outgoing connections', false]),
        Opt::SSLVersion,
        OptEnum.new('SSLVerifyMode',  [ false, 'SSL verification method', 'PEER', %W{CLIENT_ONCE FAIL_IF_NO_PEER_CERT NONE PEER}]),
        OptString.new('SSLCipher',    [ false, 'String for SSL cipher - "DHE-RSA-AES256-SHA" or "ADH"']),
        Opt::Proxies,
        Opt::CPORT,
        Opt::CHOST,
        OptInt.new('ConnectTimeout', [ true, 'Maximum number of seconds to establish a TCP connection', 10])
      ], Msf::Exploit::Remote::Tcp
    )

    register_evasion_options(
      [
        OptInt.new('TCP::max_send_size', [false, 'Maxiumum tcp segment size.  (0 = disable)', 0]),
        OptInt.new('TCP::send_delay', [false, 'Delays inserted before every send.  (0 = disable)', 0])
      ], Msf::Exploit::Remote::Tcp
    )
  end

  #
  # Establishes a TCP connection to the specified RHOST/RPORT
  #
  # @see Rex::Socket::Tcp
  # @see Rex::Socket::Tcp.create
  def connect(global = true, opts={})

    dossl = false
    if(opts.has_key?('SSL'))
      dossl = opts['SSL']
    else
      dossl = ssl
      if (datastore.default?('SSL') and rport.to_i == 443)
        dossl = true
      end
    end

    nsock = Rex::Socket::Tcp.create(
      'PeerHost'      =>  opts['RHOST'] || rhost,
      'PeerPort'      => (opts['RPORT'] || rport).to_i,
      'LocalHost'     =>  opts['CHOST'] || chost || "0.0.0.0",
      'LocalPort'     => (opts['CPORT'] || cport || 0).to_i,
      'SSL'           =>  dossl,
      'SSLVersion'    =>  opts['SSLVersion'] || ssl_version,
      'SSLVerifyMode' =>  opts['SSLVerifyMode'] || ssl_verify_mode,
      'SSLCipher'     =>  opts['SSLCipher'] || ssl_cipher,
      'Proxies'       => proxies,
      'Timeout'       => (opts['ConnectTimeout'] || connect_timeout || 10).to_i,
      'Context'       =>
        {
          'Msf'        => framework,
          'MsfExploit' => self,
        })

    # enable evasions on this socket
    set_tcp_evasions(nsock)

    # Set this socket to the global socket as necessary
    self.sock = nsock if (global)

    # Add this socket to the list of sockets created by this exploit
    add_socket(nsock)

    return nsock
  end

  # Enable evasions on a given client
  def set_tcp_evasions(socket)

    if( datastore['TCP::max_send_size'].to_i == 0 and datastore['TCP::send_delay'].to_i == 0)
      return
    end

    return if socket.respond_to?('evasive')

    socket.extend(EvasiveTCP)

    if ( datastore['TCP::max_send_size'].to_i > 0)
      socket._send_size = datastore['TCP::max_send_size']
      socket.denagle
      socket.evasive = true
    end

    if ( datastore['TCP::send_delay'].to_i > 0)
      socket._send_delay = datastore['TCP::send_delay']
      socket.evasive = true
    end
  end

  def handler(nsock = self.sock)
    # If the handler claims the socket, then we don't want it to get closed
    # during cleanup
    if ((rv = super) == Handler::Claimed)
      if (nsock == self.sock)
        self.sock = nil
      end

      # Remove this socket from the list of sockets so that it will not be
      # aborted.
      remove_socket(nsock)
    end

    return rv
  end

  #
  # Closes the TCP connection
  #
  def disconnect(nsock = self.sock)
    begin
      if (nsock)
        nsock.shutdown
        nsock.close
      end
    rescue IOError
    end

    if (nsock == sock)
      self.sock = nil
    end

    # Remove this socket from the list of sockets created by this exploit
    remove_socket(nsock)
  end

  #
  # Performs cleanup, disconnects the socket if necessary
  #
  def cleanup
    super
    disconnect
  end

  def print_prefix
    if rhost
      super + peer + " - "
    else
      super
    end
  end

  ##
  #
  # Wrappers for getters
  #
  ##

  #
  # Returns the local host for outgoing connections
  #
  def chost
    datastore['CHOST']
  end

  #
  # Returns the TCP connection timeout
  #
  def connect_timeout
    datastore['ConnectTimeout']
  end

  #
  # Returns the local port for outgoing connections
  #
  def cport
    datastore['CPORT']
  end

  #
  # Returns the local host
  #
  def lhost
    datastore['LHOST']
  end

  #
  # Returns the local port
  #
  def lport
    datastore['LPORT']
  end

  # Returns the rhost:rport
  def peer
    "#{rhost}:#{rport}"
  end

  #
  # Returns the proxy configuration
  #
  def proxies
    datastore['Proxies']
  end

  #
  # Returns the target host
  #
  def rhost
    datastore['RHOST']
  end

  #
  # Returns the remote port
  #
  def rport
    datastore['RPORT']
  end

  #
  # Returns the boolean indicating SSL
  #
  def ssl
    datastore['SSL']
  end

  #
  # Returns the string indicating SSLVersion
  #
  def ssl_version
    datastore['SSLVersion']
  end

  #
  # Returns the SSL certification verification mechanism
  #
  def ssl_verify_mode
    datastore['SSLVerifyMode']
  end

  #
  # Returns the SSL cipher to use for the context
  #
  def ssl_cipher
    datastore['SSLCipher']
  end

protected

  attr_accessor :sock

end

end