# Encoding: ASCII-8BIT

module Msf
  class Exploit
    class Remote
      # Adapted from https://github.com/rbowes-r7/libneptune
      module Unirpc
        class UniRPCError < StandardError; end
        class UniRPCCommunicationError < UniRPCError; end
        class UniRPCUnexpectedResponseError < UniRPCError; end

        # This exception is caused by using illegal values in the module and
        # probably doesn't need to be caught
        class UniRPCUsageError < UniRPCError; end

        # Argument types
        UNIRPC_TYPE_INTEGER = 0
        UNIRPC_TYPE_FLOAT = 1
        UNIRPC_TYPE_STRING = 2
        UNIRPC_TYPE_BYTES = 3

        # Message types
        UNIRPC_MESSAGE_LOGIN = 0x0F
        UNIRPC_MESSAGE_OSCOMMAND = 0x06

        def initialize(info = {})
          super

          @error_codes = YAML.safe_load(::File.join(Msf::Config.data_directory, 'unirpc-errors.yaml'))

          # This will let the module decide whether or not to use the
          # packet-level encoding
          register_advanced_options([
            OptBool.new('UNIRPC_ENCODE_MESSAGES', [true, "Use UniRPC's message encoding (which obscures messages by XORing with a constant", true])
          ])
        end

        def unirpc_get_version
          # These are the services we've found that return version numbers
          ['defcs', 'udserver'].each do |service|
            vprint_status("Trying to get version number from service #{service}...")
            connect

            sock.put(build_unirpc_message(args: [
              # Service name
              { type: :string, value: service },

              # "Secure" flag - this must be non-zero if the server is started in
              # "secure" mode (-s) - it makes no actual difference to us,
              # so just use secure mode to cover all bases
              { type: :integer, value: 1 },
            ]))

            result = recv_unirpc_message(sock)

            if result&.dig(:args, 0, :type) == :string
              version = result.dig(:args, 0, :value)&.gsub(/.*:/, '')

              unless version.nil?
                return version
              end
            end
          ensure
            disconnect
          end

          raise(UniRPCUnexpectedResponseError, 'Could not determine UniRPC version!')
        end

        # Build a unirpc packet. There are lots of arguments defined, pretty much all
        # of them optional.
        #
        # Header fields:.
        # * version_byte: The protocol version (this is always 0x6c in the protocol)
        # * other_version_byte: Another version byte (always 0x01 in the protocol)
        # * body_length_override: The length of the body (automatically calculated, normally)
        # * argcount_override: If set, specifies a custom number of "args"
        #   (automatically calculated, normally)
        #
        # Body fields:
        #
        # * body_override: If set, use it as the literal body and ignore the rest of these
        # * oldschool_data: The service supports two different types of serialized
        #   data; AFAICT, this field is just free-form string data that nothing really
        #   seems to support
        # * args: An array of arguments (the most common way to pass arguments to an
        #   rpc call).
        #
        # Args are an array of hashes with :type / :value
        # Valid types:
        # :integer - :value is the integer (32-bits)
        # :string / :bytes - value is the string or nil
        # :float - :value is just a 64-bit value
        #
        # Integer and Float values also have an :extra field, which is sent
        # where the string's length would go - I think it's normally set to
        # uninitialized memory, so probably you never need it.
        #
        # String values have a boolean :null_terminate field as well, in case
        # you want to disable null-termination (the service uses the length
        # field in some cases, and null termination in others, so it could be
        # interesting)
        #
        # Set :skip_header to not attach a header (some services require only
        # a body)
        def build_unirpc_message(
          version_byte: 0x6c,
          other_version_byte: 0x01,
          body_length_override: nil,

          argcount_override: nil,

          body_override: nil,
          oldschool_data: '',
          args: [],

          skip_header: false
        )
          encrypt = datastore['UNIRPC_ENCODE_MESSAGES']

          # Ensure this is a string (in case the caller sets it to nil or something
          oldschool_data = oldschool_data.to_s

          # Allow the caller to override the body entirely, instead of packing
          # arguments
          if body_override
            body = body_override
          else
            # Pack the args at the start of the body - this is kinda metadata-ish
            body = args.map do |a|
              case a[:type]
              when :integer
                # Ints ignore the first value, and the second is always 0
                [a[:extra] || 0, UNIRPC_TYPE_INTEGER].pack('NN')
              when :string
                # Strings store the length in the first value, and the value in the body
                if a[:null_terminate].nil? || a[:null_terminate] == true
                  [a[:value].length + 1, UNIRPC_TYPE_STRING].pack('NN')
                else
                  [a[:value].length, UNIRPC_TYPE_STRING].pack('NN')
                end
              when :bytes
                # Bytes / rpcstrings store the length in the first value, and the value in the body
                [a[:value].length, UNIRPC_TYPE_BYTES].pack('NN')
              when :float
                # Floats ignore the first value, and the second value is the type
                [a[:extra] || 0, UNIRPC_TYPE_FLOAT].pack('NN')
              else
                raise(UniRPCUsageError, "Tried to build UniRPC packet with unknown type: #{a[:type]}")
              end
            end.join

            # Follow it with the 'oldschool_data' arg
            body += oldschool_data

            # Follow that data section with the args - this is the value of the args
            body += args.map do |a|
              case a[:type]
              when :integer
                [a[:value]].pack('N')
              when :string
                str = a[:value]

                if a[:null_terminate].nil? || a[:null_terminate] == true
                  str += "\0"
                end

                # Align to multiple of 4, always adding at least one
                str += "\0"
                str += "\0" while (str.length % 4) != 0

                str
              when :bytes
                str = a[:value]

                # Alignment
                str += "\0" while (str.length % 4) != 0
                str
              when :float
                [a[:value]].pack('Q')
              else
                raise(UniRPCUsageError, "Tried to build UniRPC packet with unknown type: #{a[:type]}")
              end
            end.join
          end

          # "Encrypt" if we're supposed to
          # We use the key "2", other options include "1"
          if encrypt
            body = body.bytes.map do |b|
              (b ^ 2).chr
            end.join
          end

          # Figure out the argcount
          if argcount_override
            argcount = argcount_override
          else
            argcount = args.length

            # If we pass plaintext data, it actually counts as an extra arg
            if oldschool_data != ''
              argcount += 1
            end
          end

          # Let the user to skip appending a header, if they choose
          if skip_header
            return body
          end

          # Pack the header
          header = [
            version_byte, # Has to be 0x6c
            other_version_byte, # Can be 0x01 or 0x02
            0x00, # Reserved (ignored)
            0x00, # Reserved (ignored)

            body_length_override || body.length, # Length of data (0x7FFFFFFF => heap overflow)

            0x00000000, # Reserved (ignored)

            2,                         # Encryption "key" - basically the XOR key (can only be 1 or 2)
            0,                         # Do compression?
            encrypt ? 1 : 0,           # Encryption (0 = not encrypted, 1 = encrypted)
            0x00,                      # Padding

            0x00000000,                # Unknown (reserved?) 0 unused, but has to be 0

            argcount,                  # Argcount, which we compute earlier
            oldschool_data.length      # Data length
          ].pack('CCCCNNCCCCNnn')

          return header + body
        end

        # Receive and parse a message from UniRPC server on the given socket
        #
        # Many RPC replies put a status / error code in the first argument. To
        # check that argument and raise an error when the server returns an
        # error, set first_result_is_status to true
        def recv_unirpc_message(sock, first_result_is_status: false)
          # Receive the header
          header = sock.get_once(0x18)

          # Make sure we received all of it
          if header.nil?
            raise(UniRPCCommunicationError, "Couldn't receive UniRPC packet header")
          elsif header.length < 0x18
            raise(UniRPCCommunicationError, "UniRPC packet header was truncated (expected 24 bytes, received #{header.length}) - this might not be a UniRPC server")
          end

          # Parse out the fields
          (
            version_byte,
            other_version_byte,
            _reserved1,
            _reserved2,

            body_length,

            _reserved3,

            encryption_key,
            claim_compression,
            claim_encryption,
            _reserved4,

            _reserved5,

            argcount,
            data_length,
          ) = header.unpack('CCCCNNCCCCNnn')

          # Note that we don't attempt to decrypt / decompress here, because
          # we've never seen a server actually enable encryption or compression
          # (even if we start it)
          results = {
            header: header,
            version_byte: version_byte,
            other_version_byte: other_version_byte,
            body_length: body_length,
            encryption_key: encryption_key,
            claim_compression: claim_compression,
            claim_encryption: claim_encryption,
            argcount: argcount,
            data_length: data_length
          }

          # Receive the body
          body = sock.get_once(body_length)

          if body.length != body_length
            raise(UniRPCCommunicationError, "UniRPC packet body was truncated (expected #{body_length} bytes, received #{body.length}) - this might not be a UniRPC server")
          end

          # Parse the argument metadata, data, and argument data
          args, _data, extra_data = body.unpack("a#{argcount * 8}a#{data_length}a*")

          # Parse the argument metadata + data
          results[:args] = []
          1.upto(argcount) do
            arg, args = args.unpack('a8a*')
            (value, type) = arg.unpack('NN')

            case type
            when UNIRPC_TYPE_INTEGER # 32-bit integer
              (arg_data, extra_data) = extra_data.unpack('Na*')

              results[:args] << {
                type: :integer,
                value: arg_data,
                extra: value
              }
            when UNIRPC_TYPE_STRING # Null-able string
              if value == 0
                string_value = nil
              else
                (string, extra_data) = extra_data.unpack("a#{value}a*")
                string_value = string
              end

              results[:args] << {
                type: :string,
                value: string_value,
                extra: value
              }
            when UNIRPC_TYPE_BYTES # They call this "RPC String"
              (string, extra_data) = extra_data.unpack("a#{value}a*")
              string_value = string

              results[:args] << {
                type: :string,
                value: string_value
              }
            else
              raise(UniRPCUnexpectedResponseError, "Unidata: received unknown RPC type (#{type})!")
            end
          end

          if first_result_is_status
            if results&.dig(:args, 0, :type) != :integer
              raise(UniRPCUnexpectedResponseError, 'UniRPC server returned a non-integer status code')
            end

            error_code = results[:args][0][:value]
            if error_code != 0
              raise(UniRPCUnexpectedResponseError, "UniRPC server returned an error code: #{@error_codes[error_code] || "Unknown error: #{error_code}"}")
            end
          end

          return results
        end
      end
    end
  end
end