# -*- coding: binary -*-

module Msf
  module Exploit::Remote::SMB
    # This mixin provides a minimal SMB server
    module Server
      include Msf::Exploit::Remote::TcpServer
      include Msf::Exploit::NTLM
      CONST = ::Rex::Proto::SMB::Constants
      CRYPT = ::Rex::Proto::SMB::Crypt
      UTILS = ::Rex::Proto::SMB::Utils
      XCEPT = ::Rex::Proto::SMB::Exceptions
      EVADE = ::Rex::Proto::SMB::Evasions

      def initialize(info = {})
        super

        deregister_options('SSL', 'SSLCert')
        register_options(
          [
            OptPort.new('SRVPORT', [ true, 'The local port to listen on.', 445 ])
          ], self.class)

        register_advanced_options(
          [
            OptInt.new('SMBServerMaximumBuffer', [ true, 'The maximum number of data in megabytes to buffer', 2 ]),
            OptInt.new('SMBServerIdleTimeout', [ true, 'The maximum amount of time to keep an idle session open in seconds', 120 ])
          ], self.class)

        @smb_server_last_pool_sweep = Time.now.to_f
        @smb_server_pool_mutex = Mutex.new
        @smb_server_request_counter = 0
      end

      def setup
        super
        @state = {}
      end

      def on_client_connect(client)
        # print_status("New SMB connection from #{client.peerhost}:#{client.peerport}")
        smb_conn(client)
      end

      def on_client_data(client)
        # print_status("New data from #{client.peerhost}:#{client.peerport}")
        smb_recv(client)
        true
      end

      def on_client_close(client)
        smb_stop(client)
      end

      def smb_conn(c)
        @state[c] = {:name => "#{c.peerhost}:#{c.peerport}", :ip => c.peerhost, :port => c.peerport}
        smb_pool_update(c)
      end

      def smb_stop(c)
        # Make sure the socket is closed
        begin
          c.close
            # Handle any number of errors that a double-close or failed shutdown can trigger
        rescue ::IOError, ::EOFError,
          ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED,
          ::Errno::ETIMEDOUT, ::Errno::ENETRESET, ::Errno::ESHUTDOWN
        end

        # Delete the state table entry
        @state.delete(c)
      end

      def smb_recv(c)
        smb = @state[c]
        smb[:data] ||= ''

        buff = ''
        begin
          buff = c.get_once(-1, 0.25)
            # Handle any number of errors that a read can trigger depending on socket state
        rescue ::IOError, ::EOFError,
          ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED,
          ::Errno::ETIMEDOUT, ::Errno::ENETRESET, ::Errno::ESHUTDOWN
          vprint_status("Dropping connection from #{smb[:name]} due to exception: #{$!.class} #{$!}")
          smb_stop(c)
          return
        end

        # The client said it had data, but lied, kill the session
        unless buff and buff.length > 0
          vprint_status("Dropping connection from #{smb[:name]} due to empty payload...")
          smb_stop(c)
          return
        end

        # Append the new data to the buffer
        smb[:data] << buff

        # Prevent a simplistic DoS if the buffer is too big
        if smb[:data].length > (1024*1024*datastore['SMBServerMaximumBuffer'])
          vprint_status("Dropping connection from #{smb[:name]} due to oversized buffer of #{smb[:data].length} bytes...")
          smb_stop(c)
          return
        end

        # Update the last-seen timestamp and purge old entries
        smb_pool_update(c)

        while(smb[:data].length > 0)

          return if smb[:data].length < 4

          plen = smb[:data][2,2].unpack('n')[0]

          return if smb[:data].length < plen+4

          buff = smb[:data].slice!(0, plen+4)

          pkt_nbs = CONST::NBRAW_PKT.make_struct
          pkt_nbs.from_s(buff)

          # print_status("NetBIOS request from #{smb[:name]} #{pkt_nbs.v['Type']} #{pkt_nbs.v['Flags']} #{buff.inspect}")

          # Check for a NetBIOS name request
          if (pkt_nbs.v['Type'] == 0x81)
            # Accept any name they happen to send

            host_dst = UTILS.nbname_decode(pkt_nbs.v['Payload'][1,32]).gsub(/[\x00\x20]+$/n, '')
            host_src = UTILS.nbname_decode(pkt_nbs.v['Payload'][35,32]).gsub(/[\x00\x20]+$/n, '')

            smb[:nbdst] = host_dst
            smb[:nbsrc] = host_src

            # print_status("NetBIOS session request from #{smb[:name]} (asking for #{host_dst} from #{host_src})")
            c.write("\x82\x00\x00\x00")
            next
          end


          #
          # TODO: Support AndX parameters
          #


          # Cast this to a generic SMB structure
          pkt = CONST::SMB_BASE_PKT.make_struct
          pkt.from_s(buff)

          # Only respond to requests, ignore server replies
          if (pkt['Payload']['SMB'].v['Flags1'] & 128 != 0)
            vprint_status("Dropping connection from #{smb[:name]} due to missing client request flag")
            smb_stop(c)
            return
          end

          cmd = pkt['Payload']['SMB'].v['Command']
          begin
            smb_cmd_dispatch(cmd, c, buff)
          rescue ::Interrupt
            raise $!
          rescue ::Exception => e
            print_status("Error processing request from #{smb[:name]} (#{cmd}): #{e.class} #{e} #{e.backtrace}")
            next
          end
        end
      end

      def smb_cmd_dispatch(cmd, c, buff)
        smb = @state[c]
        print_status("Received command #{cmd} from #{smb[:name]}")
      end

      def smb_set_defaults(c, pkt)
        smb = @state[c]
        pkt['Payload']['SMB'].v['ProcessID'] = smb[:process_id].to_i
        pkt['Payload']['SMB'].v['UserID'] = smb[:user_id].to_i
        pkt['Payload']['SMB'].v['TreeID'] = smb[:tree_id].to_i
        pkt['Payload']['SMB'].v['MultiplexID'] = smb[:multiplex_id].to_i
      end

      def smb_error(cmd, c, errorclass, esn = false)
        # 0xc0000022 = Deny
        # 0xc000006D = Logon_Failure
        # 0x00000000 = Ignore
        pkt = CONST::SMB_BASE_PKT.make_struct
        smb_set_defaults(c, pkt)
        pkt['Payload']['SMB'].v['Command'] = cmd
        pkt['Payload']['SMB'].v['Flags1']  = CONST::FLAGS_REQ_RES | CONST::FLAGS_CASE_SENSITIVE
        if esn
          pkt['Payload']['SMB'].v['Flags2']  =
            CONST::FLAGS2_UNICODE_STRINGS +
            CONST::FLAGS2_EXTENDED_SECURITY +
            CONST::FLAGS2_32_BIT_ERROR_CODES +
            CONST::FLAGS2_LONG_PATH_COMPONENTS
        else
          pkt['Payload']['SMB'].v['Flags2']  =
            CONST::FLAGS2_UNICODE_STRINGS +
            CONST::FLAGS2_32_BIT_ERROR_CODES +
            CONST::FLAGS2_LONG_PATH_COMPONENTS
        end
        pkt['Payload']['SMB'].v['ErrorClass'] = errorclass
        c.put(pkt.to_s)
      end

      # Update the last-seen timestamp and purge old entries
      def smb_pool_update(c)

        @state[c][:last_action] = Time.now.to_f
        @smb_server_request_counter += 1

        unless @smb_server_request_counter % 100 == 0 ||
          @smb_server_last_pool_sweep + datastore['SMBServerIdleTimeout'].to_f < Time.now.to_f
          return
        end

        # Synchronize pool sweeps in case we move to threaded services
        @smb_server_pool_mutex.synchronize do
          purge_list = []

          @smb_server_last_pool_sweep = Time.now.to_f

          @state.keys.each do |sc|
            if @state[sc][:last_action] + datastore['SMBServerIdleTimeout'].to_f < Time.now.to_f
              purge_list << sc
            end
          end

          # Purge any idle connections to rescue file descriptors
          purge_list.each do |sc|
            vprint_status("Dropping connection from #{@state[sc][:name]} due to idle timeout...")
            smb_stop(sc)
          end
        end
      end

    end
  end

end