# -*- coding: binary -*-

module Msf

###
#
# This mixin provides a generic interface for running a socket server of some
# sort that is designed to exploit clients.  Exploits that include this mixin
# automatically take a passive stance.
#
###

module Exploit::Remote::SocketServer

  def initialize(info = {})
    super(update_info(info,
      'Stance' => Msf::Exploit::Stance::Passive))

    register_options(
      [
        OptAddressLocal.new('SRVHOST', [ true, 'The local host or network interface to listen on. This must be an address on the local machine or 0.0.0.0 to listen on all addresses.', '0.0.0.0' ]),
        OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 8080 ]),

      ], Msf::Exploit::Remote::SocketServer
    )

    register_advanced_options(
      [
        OptAddress.new('ListenerBindAddress', [ false, 'The specific IP address to bind to if different from SRVHOST']),
        OptPort.new('ListenerBindPort', [false, 'The port to bind to if different from SRVPORT']),
        OptString.new('ListenerComm', [ false, 'The specific communication channel to use for this service'])
      ], Msf::Exploit::Remote::SocketServer
    )
  end

  #
  # This mixin overrides the exploit method so that it can initiate the
  # service that corresponds with what the client has requested.
  #
  def exploit

    start_service()
    print_status("Server started.")

    # Call the exploit primer
    primer

    # Wait on the service to stop
    self.service.wait if service
  end

  #
  # Primer method to call after starting service but before handling connections
  #
  def primer
  end

  #
  # Stops the service, if one was created.
  #
  def cleanup
    super
    if service
      stopped = cleanup_service
      if stopped
        print_status("Server stopped.")
      end
    end
  end

  #
  # Called when a client has data available for reading.
  #
  def on_client_data(client)
  end

  #
  # Starts the service. Override this method in consumers
  #
  def start_service(opts = {})
  end

  #
  # Cleans up the service; either closing the socket, or deferencing the service
  #
  def cleanup_service
    if service
      begin
        if self.service.kind_of?(Rex::Service)
          temp_service = self.service
          self.service = nil
          return temp_service.deref
        else
          if self.service.kind_of?(Rex::Socket) || self.service.kind_of?(Rex::Post::Meterpreter::Channel)
            self.service.close
            self.service.stop
          end
        end

        self.service = nil
        true
      rescue ::Exception => e
        print_error(e.message)
        false
      end
    end
  end

  #
  # Returns the local host that is being listened on.
  #
  def srvhost
    datastore['SRVHOST']
  end

  #
  # Returns the local port that is being listened on.
  #
  def srvport
    datastore['SRVPORT']
  end

  #
  # Returns the address that the service is bound to. Can be different from SRVHOST when the ListenerBindAddress is
  # specified and can be used for binding to a specific address when NATing is in place.
  #
  def bindhost
    datastore['ListenerBindAddress'].blank? ? datastore['SRVHOST'] : datastore['ListenerBindAddress']
  end

  def bindport
    datastore['ListenerBindPort'].blank? ? datastore['SRVPORT'] : datastore['ListenerBindPort']
  end

  #
  # Re-generates the payload, substituting the current RHOST and RPORT with
  # the supplied client host and port from the socket.
  #
  def regenerate_payload(cli, arch = nil, platform = nil, target = nil)

    ohost = datastore['RHOST']
    oport = datastore['RPORT']
    p = nil

    begin
      # Update the datastore with the supplied client peerhost/peerport
      datastore['RHOST'] = cli.peerhost
      datastore['RPORT'] = cli.peerport

      if ((p = super(arch, platform, target)) == nil)
        print_error("Failed to generate payload")
        return nil
      end

      # Allow the payload to start a new handler
      add_handler({
        'RHOST' => datastore['RHOST'],
        'RPORT' => datastore['RPORT']
      })

    ensure
      datastore['RHOST'] = ohost
      datastore['RPORT'] = oport
    end

    p
  end

protected

  #
  # Determines appropriate listener comm
  #
  def _determine_server_comm(ip, srv_comm = datastore['ListenerComm'].to_s)
    comm = nil

    case srv_comm
    when 'local'
      comm = ::Rex::Socket::Comm::Local
    when /\A-?[0-9]+\Z/
      comm = framework.sessions.get(srv_comm.to_i)
      raise(RuntimeError, "Socket Server Comm (Session #{srv_comm}) does not exist") unless comm
      raise(RuntimeError, "Socket Server Comm (Session #{srv_comm}) does not implement Rex::Socket::Comm") unless comm.is_a? ::Rex::Socket::Comm
    when nil, ''
      unless ip.nil?
        comm = Rex::Socket::SwitchBoard.best_comm(ip)
      end
    else
      raise(RuntimeError, "SocketServer Comm '#{srv_comm}' is invalid")
    end

    comm || ::Rex::Socket::Comm::Local
  end

  def via_string(comm)
    if comm.respond_to?(:type) && comm.respond_to?(:sid)
      via = "via the #{comm.type} on session #{comm.sid}"
    else
      via = ""
    end

    via
  end

  attr_accessor :service # :nodoc:

end

end
