# -*- 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(
      [
        OptAddress.new('SRVHOST', [ true, "The local host to listen on. This must be an address on the local machine or 0.0.0.0", '0.0.0.0' ]),
        OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 8080 ]),

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

    register_advanced_options(
      [
        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
  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)
      stop_service()
      print_status("Server stopped.")
    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(*args)
  end

  #
  # Stops the service.
  #
  def stop_service
    if (service)
      begin
        self.service.deref if self.service.kind_of?(Rex::Service)
        if self.service.kind_of?(Rex::Socket) || self.service.kind_of?(Rex::Post::Meterpreter::Channel)
          self.service.close
          self.service.stop
        end

        self.service = nil
      rescue ::Exception
      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

  #
  # 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(srv_comm = datastore['ListenerComm'].to_s)
    case srv_comm
    when 'local'
      comm = ::Rex::Socket::Comm::Local
    when /\A[0-9]+\Z/
      comm = framework.sessions[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, ''
      comm = nil 
    else
      raise(RuntimeError, "SocketServer Comm '#{srv_comm}' is invalid")
    end 

    comm
  end

  def via_string_for_ip(ip, comm)
    comm_used = comm
    comm_used ||= Rex::Socket::SwitchBoard.best_comm(ip)
    comm_used ||= Rex::Socket::Comm::Local

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

    via
  end

  attr_accessor :service # :nodoc:

end

end