# -*- coding: binary -*-
require 'rex/proto/rmi'
require 'rex/java/serialization'
require 'stringio'
require 'msf/core/exploit/java/rmi/util'
require 'msf/core/exploit/java/rmi/builder'
require 'msf/core/exploit/java/rmi/client/registry'
require 'msf/core/exploit/java/rmi/client/jmx'

module Msf
  class Exploit
    class Remote
      module Java
        module Rmi
          module Client

            include Msf::Exploit::Remote::Java::Rmi::Util
            include Msf::Exploit::Remote::Java::Rmi::Builder
            include Msf::Exploit::Remote::Java::Rmi::Client::Registry
            include Msf::Exploit::Remote::Java::Rmi::Client::Jmx
            include Msf::Exploit::Remote::Tcp

            def initialize(info = {})
              super

              register_advanced_options(
                [
                  OptInt.new('RmiReadLoopTimeout', [ true, 'Maximum number of seconds to wait for data between read iterations', 1])
                ], Msf::Exploit::Remote::Java::Rmi::Client
              )
            end

            # Returns the timeout to wait for data between read iterations
            #
            # @return [Integer]
            def read_loop_timeout
              datastore['RmiReadLoopTimeout'] || 1
            end

            # Returns the target host
            #
            # @return [String]
            def rhost
              datastore['RHOST']
            end

            # Returns the target port
            #
            # @return [Integer]
            def rport
              datastore['RPORT']
            end


            # Sends a RMI header stream
            #
            # @param opts [Hash]
            # @option opts [Rex::Socket::Tcp] :sock
            # @return [Integer] the number of bytes sent
            # @see Msf::Rmi::Client::Streams#build_header
            def send_header(opts = {})
              nsock = opts[:sock] || sock
              stream = build_header(opts)
              nsock.put(stream.encode + "\x00\x00\x00\x00\x00\x00")
            end

            # Sends a RMI CALL stream
            #
            # @param opts [Hash]
            # @option opts [Rex::Socket::Tcp] :sock
            # @option opts [Rex::Proto::Rmi::Model::Call] :call
            # @return [Integer] the number of bytes sent
            # @see Msf::Rmi::Client::Streams#build_call
            def send_call(opts = {})
              nsock = opts[:sock] || sock
              call = opts[:call] || build_call(opts)
              nsock.put(call.encode)
            end

            # Sends a RMI DGCACK stream
            #
            # @param opts [Hash]
            # @option opts [Rex::Socket::Tcp] :sock
            # @return [Integer] the number of bytes sent
            # @see Msf::Rmi::Client::Streams#build_dgc_ack
            def send_dgc_ack(opts = {})
              nsock = opts[:sock] || sock
              stream = build_dgc_ack(opts)
              nsock.put(stream.encode)
            end

            # Reads the Protocol Ack
            #
            # @param opts [Hash]
            # @option opts [Rex::Socket::Tcp] :sock
            # @return [Rex::Proto::Rmi::Model::ProtocolAck] if success
            # @return [NilClass] otherwise
            # @see Rex::Proto::Rmi::Model::ProtocolAck.decode
            def recv_protocol_ack(opts = {})
              nsock = opts[:sock] || sock
              data = safe_get_once(nsock)
              begin
                ack = Rex::Proto::Rmi::Model::ProtocolAck.decode(StringIO.new(data))
              rescue Rex::Proto::Rmi::DecodeError
                return nil
              end

              ack
            end

            # Reads a ReturnData message and returns the java serialized stream
            # with the return data value.
            #
            # @param opts [Hash]
            # @option opts [Rex::Socket::Tcp] :sock
            # @return [Rex::Proto::Rmi::Model::ReturnValue] if success
            # @return [NilClass] otherwise
            # @see Rex::Proto::Rmi::Model::ReturnData.decode
            def recv_return(opts = {})
              nsock = opts[:sock] || sock
              data = safe_get_once(nsock)

              begin
                return_data = Rex::Proto::Rmi::Model::ReturnData.decode(StringIO.new(data))
              rescue Rex::Proto::Rmi::DecodeError
                return nil
              end

              return_data.return_value
            end

            # Helper method to read fragmented data from a ```Rex::Socket::Tcp```
            #
            # @param nsock [Rex::Socket::Tcp]
            # @return [String]
            def safe_get_once(nsock = sock, loop_timeout = read_loop_timeout)
              data = ''
              begin
                res = nsock.get_once
              rescue ::EOFError
                res = nil
              end

              while res && nsock.has_read_data?(loop_timeout)
                data << res
                begin
                  res = nsock.get_once
                rescue ::EOFError
                  res = nil
                end
              end

              data << res if res
              data
            end
          end

        end
      end
    end
  end
end
