# -*- coding: binary -*-
require 'rex/socket'
require 'thread'

module Msf
module Handler

###
#
# This module implements the reverse TCP handler.  This means
# that it listens on a port waiting for a connection until
# either one is established or it is told to abort.
#
# This handler depends on having a local host and port to
# listen on.
#
###
module ReverseTcp

  include Msf::Handler
  include Msf::Handler::Reverse
  include Msf::Handler::Reverse::Comm

  #
  # Returns the string representation of the handler type, in this case
  # 'reverse_tcp'.
  #
  def self.handler_type
    return "reverse_tcp"
  end

  #
  # Returns the connection-described general handler type, in this case
  # 'reverse'.
  #
  def self.general_handler_type
    "reverse"
  end

  #
  # Initializes the reverse TCP handler and ads the options that are required
  # for all reverse TCP payloads, like local host and local port.
  #
  def initialize(info = {})
    super

    # XXX: Not supported by all modules
    register_advanced_options(
      [
        OptInt.new('ReverseConnectRetries', [ true, 'The number of connection attempts to try before exiting the process', 5 ]),
        OptAddress.new('ReverseListenerBindAddress', [ false, 'The specific IP address to bind to on the local system']),
        OptBool.new('ReverseListenerThreaded', [ true, 'Handle every connection in a new thread (experimental)', false])
      ], Msf::Handler::ReverseTcp)

    self.conn_threads = []
  end


  #
  # Closes the listener socket if one was created.
  #
  def cleanup_handler
    stop_handler

    # Kill any remaining handle_connection threads that might
    # be hanging around
    conn_threads.each { |thr|
      thr.kill rescue nil
    }
  end

  # A string suitable for displaying to the user
  #
  # @return [String]
  def human_name
    "reverse TCP"
  end

  #
  # Starts monitoring for an inbound connection.
  #
  def start_handler
    queue = ::Queue.new

    local_port = bind_port

    self.listener_thread = framework.threads.spawn("ReverseTcpHandlerListener-#{local_port}", false, queue) { |lqueue|
      loop do
        # Accept a client connection
        begin
          client = self.listener_sock.accept
          if ! client
            wlog("ReverseTcpHandlerListener-#{local_port}: No client received in call to accept, exiting...")
            break
          end

          self.pending_connections += 1
          lqueue.push(client)
        rescue ::Exception
          wlog("ReverseTcpHandlerListener-#{local_port}: Exception raised during listener accept: #{$!}\n\n#{$@.join("\n")}")
          break
        end
      end
    }

    self.handler_thread = framework.threads.spawn("ReverseTcpHandlerWorker-#{local_port}", false, queue) { |cqueue|
      loop do
        begin
          client = cqueue.pop

          if ! client
            elog("ReverseTcpHandlerWorker-#{local_port}: Queue returned an empty result, exiting...")
            break
          end

          # Timeout and datastore options need to be passed through to the client
          opts = {
            :datastore    => datastore,
            :expiration   => datastore['SessionExpirationTimeout'].to_i,
            :comm_timeout => datastore['SessionCommunicationTimeout'].to_i,
            :retry_total  => datastore['SessionRetryTotal'].to_i,
            :retry_wait   => datastore['SessionRetryWait'].to_i
          }

          if datastore['ReverseListenerThreaded']
            self.conn_threads << framework.threads.spawn("ReverseTcpHandlerSession-#{local_port}-#{client.peerhost}", false, client) { |client_copy|
              handle_connection(wrap_aes_socket(client_copy), opts)
            }
          else
            handle_connection(wrap_aes_socket(client), opts)
          end
        rescue ::Exception
          elog("Exception raised from handle_connection: #{$!.class}: #{$!}\n\n#{$@.join("\n")}")
        end
      end
    }

  end

  def wrap_aes_socket(sock)
    if datastore["PAYLOAD"] !~ /java\// or (datastore["AESPassword"] || "") == ""
      return sock
    end

    socks = Rex::Socket::tcp_socket_pair()
    socks[0].extend(Rex::Socket::Tcp)
    socks[1].extend(Rex::Socket::Tcp)

    m = OpenSSL::Digest.new('md5')
    m.reset
    key = m.digest(datastore["AESPassword"] || "")

    Rex::ThreadFactory.spawn('Session-AESEncrypt', false) {
      c1 = OpenSSL::Cipher.new('aes-128-cfb8')
      c1.encrypt
      c1.key=key
      sock.put([0].pack('N'))
      sock.put(c1.iv=c1.random_iv)
      buf1 = socks[0].read(4096)
      while buf1 and buf1 != ""
        sock.put(c1.update(buf1))
        buf1 = socks[0].read(4096)
      end
      sock.close()
    }
    Rex::ThreadFactory.spawn('Session-AESDecrypt', false) {
      c2 = OpenSSL::Cipher.new('aes-128-cfb8')
      c2.decrypt
      c2.key=key
      iv=""
      while iv.length < 16
        iv << sock.read(16-iv.length)
      end
      c2.iv = iv
      buf2 = sock.read(4096)
      while buf2 and buf2 != ""
        socks[0].put(c2.update(buf2))
        buf2 = sock.read(4096)
      end
      socks[0].close()
    }
    return socks[1]
  end

  #
  # Stops monitoring for an inbound connection.
  #
  def stop_handler
    # Terminate the listener thread
    if (self.listener_thread and self.listener_thread.alive? == true)
      self.listener_thread.kill
      self.listener_thread = nil
    end

    # Terminate the handler thread
    if (self.handler_thread and self.handler_thread.alive? == true)
      self.handler_thread.kill
      self.handler_thread = nil
    end

    if (self.listener_sock)
      begin
        self.listener_sock.close
      rescue IOError
        # Ignore if it's listening on a dead session
        dlog("IOError closing listener sock; listening on dead session?", LEV_1)
      end
      self.listener_sock = nil
    end
  end

protected

  attr_accessor :listener_sock # :nodoc:
  attr_accessor :listener_thread # :nodoc:
  attr_accessor :handler_thread # :nodoc:
  attr_accessor :conn_threads # :nodoc:
end

end
end