# -*- coding: binary -*-

module Msf

###
#
# This module acts as a base for all handler pseudo-modules.  They aren't
# really modules, so don't get the wrong idea champs!  They're merely
# mixed into dynamically generated payloads to handle monitoring for
# a connection.  Handlers are layered in between the base payload
# class and any other payload class.  A super cool ASCII diagram would
# look something like this
#
#      Module
#        ^
#        |
#     Payload
#        ^
#        |
#     Handler
#        ^
#        |
#      Stager
#        ^
#        |
#       Stage
#
###
module Handler

  ##
  #
  # Constants used with the ``handler'' method to indicate whether or not the
  # connection was used.
  #
  ##

  #
  # Returned by handlers to indicate that a socket has been claimed for use
  # by the payload.
  #
  Claimed = "claimed"
  #
  # Returned by handlers to indicate that a socket has not been claimed for
  # use.
  #
  Unused  = "unused"

  #
  # Returns the handler type.
  #
  def self.handler_type
    return "none"
  end

  #
  # Returns the transport-independent handler type.
  #
  def self.general_handler_type
    "none"
  end

  #
  # Returns the handler's name, if any.
  #
  def handler_name
    module_info['HandlerName']
  end

  #
  # Initializes the session waiter event and other fun stuff.
  #
  def initialize(info = {})
    super

    # Initialize the pending_connections counter to 0
    self.pending_connections = 0

    # Initialize the sessions counter to 0
    self.sessions = 0

    # Create the waiter event with auto_reset set to false so that
    # if a session is ever created, waiting on it returns immediately.
    self.session_waiter_event = Rex::Sync::Event.new(false, false)
  end

  #
  # Sets up the connection handler.
  #
  def setup_handler
  end

  #
  # Terminates the connection handler.
  #
  def cleanup_handler
  end

  #
  # Start monitoring for a connection.
  #
  def start_handler
  end

  #
  # Start another connection monitor
  #
  def add_handler(opts={})
  end

  #
  # Stop monitoring for a connection.
  #
  def stop_handler
  end

  #
  # Checks to see if a payload connection has been established on
  # the supplied connection.  This is necessary for find-sock style
  # payloads.
  #
  def handler(sock)
  end

  #
  # Handles an established connection supplied in the in and out
  # handles.  The handles are passed as parameters in case this
  # handler is capable of handling multiple simultaneous
  # connections.  The default behavior is to attempt to create a session for
  # the payload.  This path will not be taken for multi-staged payloads.
  #
  def handle_connection(conn, opts={})
    create_session(conn, opts)
  end

  #
  # The amount of time to wait for a session to come in.
  #
  def wfs_delay
    2
  end

  #
  # Waits for a session to be created as the result of a handler connection
  # coming in.  The return value is a session object instance on success or
  # nil if the timeout expires.
  #
  def wait_for_session(t = wfs_delay)
    session = nil

    begin
      session = session_waiter_event.wait(t)
    rescue ::Timeout::Error
    end

    # If a connection has arrived, wait longer...
    if (pending_connections > 0)
      session = session_waiter_event.wait
    end

    return session
  end

  #
  # Interrupts a wait_for_session call by notifying with a nil event
  #
  def interrupt_wait_for_session
    return unless session_waiter_event
    session_waiter_event.notify(nil)
  end

  #
  # Set by the exploit module to configure handler
  #
  attr_accessor :exploit_config

  #
  # This will be non-nil if the handler has a parent payload that it
  # was spawned from.  Right now, this is only the case with generic
  # payloads.  The parent payload is used to create a session
  # rather than using the instance itself.
  #
  attr_accessor :parent_payload

protected

  #
  # Creates a session, if necessary, for the connection that's been handled.
  # Sessions are only created if the payload that's been mixed in has an
  # associated session.
  #
  def create_session(conn, opts={})
    # If there is a parent payload, then use that in preference.
    return parent_payload.create_session(conn, opts) if (parent_payload)

    # If the payload we merged in with has an associated session factory,
    # allocate a new session.
    if (self.session)
      begin
        # if there's a create_session method then use it, as this
        # can form a factory for arb session types based on the
        # payload.
        if self.session.respond_to?('create_session')
          s = self.session.create_session(conn, opts)
        else
          s = self.session.new(conn, opts)
        end
      rescue ::Exception => e
        # We just wanna show and log the error, not trying to swallow it.
        print_error("#{e.class} #{e.message}")
        elog('Could not allocate a new Session.', error: e)
        raise e
      end

      # Pass along the framework context
      s.framework = framework

      # Associate this system with the original exploit
      # and any relevant information
      s.set_from_exploit(assoc_exploit)

      # set injected workspace value if db is active
      if framework.db.active && wspace = framework.db.find_workspace(s.workspace)
        framework.db.workspace = wspace
      end

      # Pass along any associated payload uuid if specified
      if opts[:payload_uuid]
        s.payload_uuid = opts[:payload_uuid]
        s.payload_uuid.registered = false
        if framework.db.active
          payload_info = { uuid: s.payload_uuid.puid_hex, workspace: framework.db.workspace }
          uuid_info = framework.db.payloads(payload_info).first
        else
          print_warning('Without a database connected that payload UUID tracking will not work!')
        end
        if s.payload_uuid.respond_to?(:puid_hex) && uuid_info
          s.payload_uuid.registered = true
          s.payload_uuid.name = uuid_info['name']
          s.payload_uuid.timestamp = uuid_info['timestamp']
        else
          s.payload_uuid.registered = false
        end
      end

      # If the session is valid, register it with the framework and
      # notify any waiters we may have.
      if (s)
        # Defer the session registration to the Session Manager scheduler
        registration = Proc.new do
          register_session(s)
        end
        framework.sessions.schedule registration
      end

      return s
    end
    nil
  end

  #
  # Registers a session with the framework and notifies any waiters of the
  # new session.
  #
  def register_session(session)
    # Register the session with the framework
    framework.sessions.register(session)

    # Call the handler's on_session() method
    if session.respond_to?(:bootstrap)
      session.bootstrap(datastore, self)

      return unless session.alive
    end

    # Process the auto-run scripts for this session
    if session.respond_to?(:process_autoruns)
      session.process_autoruns(datastore)
    end

    # Tell the handler that we have a session
    on_session(session)

    # Notify the framework that we have a new session opening up...
    # Don't let errant event handlers kill our session
    begin
      framework.events.on_session_open(session)
    rescue ::Exception => e
      wlog("Exception in on_session_open event handler: #{e.class}: #{e}")
      wlog("Call Stack\n#{e.backtrace.join("\n")}")
    end

    # If there is an exploit associated with this payload, then let's notify
    # anyone who is interested that this exploit succeeded
    if assoc_exploit
      framework.events.on_exploit_success(assoc_exploit, session)
    end

    # Notify waiters that they should be ready to rock
    session_waiter_event.notify(session)

    # Decrement the pending connections counter now that we've processed
    # one session.
    self.pending_connections -= 1

    # Count the number of sessions we have registered
    self.sessions += 1
  end

  attr_accessor :session_waiter_event # :nodoc:
  attr_accessor :pending_connections  # :nodoc:
  attr_accessor :sessions # :nodoc:

end

end
