# -*- coding: binary -*-

module Msf

###
#
# This mixin provides a minimal SMB server
#
###

module Exploit::Remote::SMBServer
  include Exploit::Remote::TcpServer
  include 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

    register_options(
      [
        OptPort.new('SRVPORT',    [ true, "The local port to listen on.", 445 ])
      ], self.class)
  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}
  end

  def smb_stop(c)
    @state.delete(c)
  end

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

    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 response to requests, ignore server replies
      if (pkt['Payload']['SMB'].v['Flags1'] & 128 != 0)
        print_status("Ignoring server response from #{smb[:name]}")
        next
      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']  = 0x88
    if esn
      pkt['Payload']['SMB'].v['Flags2']  = 0xc801
    else
      pkt['Payload']['SMB'].v['Flags2']  = 0xc001
    end
    pkt['Payload']['SMB'].v['ErrorClass'] = errorclass
    c.put(pkt.to_s)
  end

end

end