module Msf
module Exploit::SMB
  # This class provides a Man-In-The-Middle packet Dispatcher.
  class ShadowMitmDispatcher < RubySMB::Dispatcher::Base
      READ_TIMEOUT = 30
      TCP_MSS = 536 # RFC 879

      # The underlying stream that we select on
      # @!attribute [rw] stream
      #   @return [IO]
      attr_accessor :stream

      alias :tcp_socket :stream

      # The read timeout
      # @!attribute [rw] read_timeout
      #   @return [Integer]
      attr_accessor :read_timeout

      # The mitm interface
      # @!attribute [rw] interface
      #   @return [String]
      attr_accessor :interface

      # The mac address
      # @!attribute [rw] mac
      #   @return [String]
      attr_accessor :mac

      # The source mac address
      # @!attribute [rw] eth_src
      #   @return [String]
      attr_accessor :eth_src

      # The destination mac address
      # @!attribute [rw] eth_dst
      #   @return [String]
      attr_accessor :eth_dst

      # The source ip address
      # @!attribute [rw] ip_src
      #   @return [String]
      attr_accessor :ip_src

      # The destination ip address
      # @!attribute [rw] ip_dst
      #   @return [String]
      attr_accessor :ip_dst

      # The tcp sequence number
      # @!attribute [rw] tcp_seq
      #   @return [Integer]
      attr_accessor :tcp_seq

      # The tcp acknowledgement number
      # @!attribute [rw] tcp_ack
      #   @return [Integer]
      attr_accessor :tcp_ack

      # The tcp source port
      # @!attribute [rw] tcp_src
      #   @return [Integer]
      attr_accessor :tcp_src

      # The tcp destination port
      # @!attribute [rw] tcp_dst
      #   @return [Integer]
      attr_accessor :tcp_dst

      # The tcp window size
      # @!attribute [rw] tcp_win
      #   @return [Integer]
      attr_accessor :tcp_win

      # The tcp max segment size
      # @!attribute [rw] tcp_mss
      #   @return [Integer]
      attr_accessor :tcp_mss

      # The Current Capture
      # @!attribute [rw] ccap
      #   @return [PacketFu::Capture]
      attr_accessor :ccap

      # We decide the host based on the first person to connect
      def initialize(interface: , mac: , eth_src: , eth_dst: , ip_src: , ip_dst: , tcp_src:, tcp_dst: , tcp_seq: , tcp_ack: , tcp_win: , tcp_mss: TCP_MSS, read_timeout: READ_TIMEOUT)
        @interface = interface
        @mac = mac
        @eth_src = eth_src
        @eth_dst = eth_dst
        @ip_src = ip_src
        @ip_dst = ip_dst
        @tcp_src = tcp_src
        @tcp_dst = tcp_dst
        @tcp_seq = tcp_seq
        @tcp_ack = tcp_ack
        @tcp_win = tcp_win
        @tcp_mss = tcp_mss
        @read_timeout = read_timeout
        # Just create something to close
        @stream = Socket.new Socket::AF_INET, Socket::SOCK_STREAM
      end

      # @param host [String] passed to TCPSocket.new
      # @param port [Fixnum] passed to TCPSocket.new
      def self.connect#(host, port: 445, socket: TCPSocket.new(host, port))
        new()
      end

      # Needs: @tcp_src, @tcp_dst, @tcp_seq, @tcp_ack, @tcp-win, @interface, @mac
      def send_packet(packet, nbss_header: true, tcp_flags: { ack: 1, psh: 1 }, tcp_opts: "")
        data = nbss_header ? nbss(packet) : ''
        if packet.respond_to?(:to_binary_s)
          data << packet.to_binary_s
        else
          data << packet
        end
        eth_header = PacketFu::EthHeader.new(eth_src: @eth_src, eth_dst: @eth_dst)
        ip_header = PacketFu::IPHeader.new(ip_src: @ip_src, ip_dst: @ip_dst)
        tcp_header = PacketFu::TCPHeader.new(
          tcp_src: @tcp_src,
          tcp_dst: @tcp_dst,
          tcp_seq: @tcp_seq,
          tcp_ack: @tcp_ack,
          tcp_win: @tcp_win,
          tcp_flags: tcp_flags,
          tcp_opts: tcp_opts
        )
        pkt = PacketFu::TCPPacket.new(eth: eth_header, ip: ip_header, tcp: tcp_header)
        pkt.payload = data
        pkt.recalc
        @stream.close if @stream
        @ccap = PacketFu::Capture.new(
          iface: @interface,
          promisc: false,
          start: true,
          filter: "ether dst #{@mac} and not ether src #{@mac} and src port #{@tcp_dst} and dst port #{@tcp_src}"
        )
        @stream = @ccap.stream
        if data.size > @tcp_mss
          ip_body = pkt.ip_header.body.to_s
          ip_id = pkt.ip_header.ip_id
          ip_body.chars.each_slice(@tcp_mss+40).each_with_index do |frag, index|
            fpkt = PacketFu::IPPacket.new(eth: eth_header.to_s, ip: ip_header.to_s)
            fpkt.ip_id = ip_id
            fpkt.ip_header.ip_frag = 0x20*0x100 unless (index+1)*(@tcp_mss+40) >= ip_body.size
            fpkt.ip_header.ip_frag = fpkt.ip_header.ip_frag + ((@tcp_mss+40)/8)*index
            fpkt.payload = frag.join
            fpkt.ip_header.ip_recalc(:ip_len)
            fpkt.ip_header.ip_recalc(:ip_sum)
            fpkt.to_w(@interface)
          end
        else
          pkt.to_w(@interface)
        end
        nil
      end

      # Read a packet off the wire and parse it into a string
      #
      # @param full_response [Boolean] whether to include the NetBios Session Service header in the response
      # @return [String] the raw response (including the NetBios Session Service header if full_response is true)
      # @raise [RubySMB::Error::NetBiosSessionService] if there's an error reading the first 4 bytes,
      #   which are assumed to be the NetBiosSessionService header.
      # @raise [RubySMB::Error::CommunicationError] if the read timeout expires or an error occurs when reading the packet
      def recv_packet(full_response: false)
        raise RubySMB::Error::CommunicationError, 'Capture has not been initialized' unless @ccap.class == PacketFu::Capture
        pkt = @stream.each_data do |data|
          pkt = PacketFu::Packet.parse(data)
          break pkt if (pkt.tcp_header.tcp_seq == @tcp_ack) || (@tcp_ack == 0)
        end
        raise Errno::ECONNRESET, 'Received a RST packet' if pkt.tcp_header.tcp_flags.rst == 1
        #puts "#{pkt.tcp_header.tcp_seq} == #{@tcp_ack}"
        @tcp_ack = pkt.tcp_header.tcp_seq + pkt.tcp_header.body.size
        @tcp_ack += 1 if pkt.tcp_header.tcp_flags.syn == 1
        @tcp_seq = pkt.tcp_header.tcp_ack
        return pkt.payload[(full_response ? 0 : 4)..-1]
      rescue Errno::EINVAL, Errno::ECONNABORTED, Errno::ECONNRESET, TypeError, NoMethodError => e
        raise RubySMB::Error::CommunicationError, "An error occurred reading from the Network #{e.message}"
      end
    end
  end
end
