Add support for Ping and PingAck

lib/rex/proto/rmi/model.rb

@@ -26,4 +26,6 @@ require 'rex/proto/rmi/model/protocol_ack'
 require 'rex/proto/rmi/model/continuation'
 require 'rex/proto/rmi/model/call'
 require 'rex/proto/rmi/model/return_data'
-require 'rex/proto/rmi/model/dbg_ack'
+require 'rex/proto/rmi/model/dbg_ack'
+require 'rex/proto/rmi/model/ping'
+require 'rex/proto/rmi/model/ping_ack'

lib/rex/proto/rmi/model/ping.rb

@@ -0,0 +1,41 @@
+# -*- coding: binary -*-
+
+module Rex
+  module Proto
+    module Rmi
+      module Model
+        # This class provides a representation of an RMI Ping stream. A Ping is a message for testing
+        # livereness of a remote virtual machine.
+        class Ping < Element
+
+          # @!attribute stream_id
+          #   @return [Fixnum] the input stream id
+          attr_accessor :stream_id
+
+          private
+
+          # Reads the stream id from the IO
+          #
+          # @param io [IO] the IO to read from
+          # @return [String]
+          # @raise [RuntimeError] if fails to decode stream id
+          def decode_stream_id(io)
+            stream_id = read_byte(io)
+            unless stream_id == PING_MESSAGE
+              raise ::RuntimeError, 'Failed to decode Ping stream id'
+            end
+
+            stream_id
+          end
+
+          # Encodes the stream_id field
+          #
+          # @return [String]
+          def encode_stream_id
+            [stream_id].pack('C')
+          end
+        end
+      end
+    end
+  end
+end

lib/rex/proto/rmi/model/ping_ack.rb

@@ -0,0 +1,41 @@
+# -*- coding: binary -*-
+
+module Rex
+  module Proto
+    module Rmi
+      module Model
+        # This class provides a representation of an RMI PingAck stream. A PingAck is the acknowledgement
+        # for a Ping message.
+        class PingAck < Element
+
+          # @!attribute stream_id
+          #   @return [Fixnum] the input stream id
+          attr_accessor :stream_id
+
+          private
+
+          # Reads the stream id from the IO
+          #
+          # @param io [IO] the IO to read from
+          # @return [String]
+          # @raise [RuntimeError] if fails to decode stream id
+          def decode_stream_id(io)
+            stream_id = read_byte(io)
+            unless stream_id == PING_ACK
+              raise ::RuntimeError, 'Failed to decode PingAck stream id'
+            end
+
+            stream_id
+          end
+
+          # Encodes the stream_id field
+          #
+          # @return [String]
+          def encode_stream_id
+            [stream_id].pack('C')
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/rex/proto/rmi/model/ping_ack_spec.rb

@@ -0,0 +1,37 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+require 'stringio'
+require 'rex/proto/rmi'
+
+describe Rex::Proto::Rmi::Model::PingAck do
+
+  subject(:ping_ack) do
+    described_class.new
+  end
+
+  let(:sample) do
+    "\x53"
+  end
+
+  let(:sample_io) { StringIO.new(sample) }
+
+  describe "#decode" do
+    it "returns the Rex::Proto::Rmi::Model::PingAck decoded" do
+      expect(ping_ack.decode(sample_io)).to eq(ping_ack)
+    end
+
+    it "decodes stream_id correctly" do
+      ping_ack.decode(sample_io)
+      expect(ping_ack.stream_id).to eq(Rex::Proto::Rmi::Model::PING_ACK)
+    end
+  end
+
+  describe "#encode" do
+    it "encodes the PingAck correctly" do
+      ping_ack.stream_id = Rex::Proto::Rmi::Model::PING_ACK
+      expect(ping_ack.encode).to eq(sample)
+    end
+  end
+end
+

spec/lib/rex/proto/rmi/model/ping_spec.rb

@@ -0,0 +1,36 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+require 'stringio'
+require 'rex/proto/rmi'
+
+describe Rex::Proto::Rmi::Model::Ping do
+
+  subject(:ping) do
+    described_class.new
+  end
+
+  let(:sample) do
+    "\x52"
+  end
+
+  let(:sample_io) { StringIO.new(sample) }
+
+  describe "#decode" do
+    it "returns the Rex::Proto::Rmi::Model::Ping decoded" do
+      expect(ping.decode(sample_io)).to eq(ping)
+    end
+
+    it "decodes stream_id correctly" do
+      ping.decode(sample_io)
+      expect(ping.stream_id).to eq(Rex::Proto::Rmi::Model::PING_MESSAGE)
+    end
+  end
+
+  describe "#encode" do
+    it "encodes the Ping correctly" do
+      ping.stream_id = Rex::Proto::Rmi::Model::PING_MESSAGE
+      expect(ping.encode).to eq(sample)
+    end
+  end
+end