lib/ronin/formatting/extensions/binary/integer.rb in ronin-support-0.4.1 vs lib/ronin/formatting/extensions/binary/integer.rb in ronin-support-0.5.0.rc1

@@ -15,35 +15,37 @@
 #
 # You should have received a copy of the GNU Lesser General Public License
 # along with Ronin Support.  If not, see <http://www.gnu.org/licenses/>.
 #
 
+require 'ronin/binary/template'
+
 class Integer
 
   #
   # Extracts a sequence of bytes which represent the Integer.
   #
-  # @param [Integer] address_length
+  # @param [Integer] length
   #   The number of bytes to decode from the Integer.
   #
   # @param [Symbol, String] endian
   #   The endianness to use while decoding the bytes of the Integer.
   #   May be one of:
   #
-  #   * `:big` / `"big"`
-  #   * `:little` / `"little"`
-  #   * `:net` / `"net"`
+  #   * `big`
+  #   * `little`
+  #   * `net`
   #
   # @return [Array]
   #   The bytes decoded from the Integer.
   #
   # @raise [ArgumentError]
   #   The given `endian` was not one of:
   #
-  #   * `:little` / `"little"`
-  #   * `:net` / `"net"`
-  #   * `:big` / `"big"`
+  #   * `little`
+  #   * `net`
+  #   * `big`
   #
   # @example
   #   0xff41.bytes(2)
   #   # => [65, 255]
   #
@@ -51,30 +53,30 @@
   #   0xff41.bytes(4, :big)
   #   # => [0, 0, 255, 65]
   #
   # @api public
   #
-  def bytes(address_length,endian=:little)
+  def bytes(length,endian=:little)
     endian = endian.to_sym
     buffer = []
 
     case endian
     when :little, :net
       mask  = 0xff
       shift = 0
 
-      address_length.times do |i|
+      length.times do |i|
         buffer << ((self & mask) >> shift)
 
         mask <<= 8
         shift += 8
       end
     when :big
-      shift = ((address_length - 1) * 8)
+      shift = ((length - 1) * 8)
       mask  = (0xff << shift)
 
-      address_length.times do |i|
+      length.times do |i|
         buffer << ((self & mask) >> shift)
 
         mask >>= 8
         shift -= 8
       end
@@ -84,66 +86,77 @@
 
     return buffer
   end
 
   #
-  # Packs the Integer into a String, for a specific architecture and
-  # address-length.
+  # Packs the Integer into a String.
   #
-  # @param [Ronin::Arch, #endian, #address_length, String] arch
-  #   The architecture to pack the Integer for.
+  # @param [String, Symbol, arch] arguments
+  #   The `Array#pack` code, {Ronin::Binary::Template} type or Architecture
+  #   object with `#endian` and `#address_length` methods.
   #
-  # @param [Integer] address_length
-  #   The number of bytes to pack.
-  #
   # @return [String]
   #   The packed Integer.
   #
   # @raise [ArgumentError]
-  #   The given `arch` does not respond to the `endian` or
-  #   `address_length` methods.
+  #   The arguments were not a String, Symbol or Architecture object.
   #
-  # @example using archs other than `Ronin::Arch`.
+  # @example using a `Array#pack` template:
+  #   0x41.pack('V')
+  #   # => "A\0\0\0"
+  #
+  # @example using {Ronin::Binary::Template} types:
+  #   0x41.pack(:uint32_le)
+  #
+  # @example using archs other than `Ronin::Arch` (**deprecated**):
   #   arch = OpenStruct.new(:endian => :little, :address_length => 4)
-  #   
+  #
   #   0x41.pack(arch)
   #   # => "A\0\0\0"
   #
-  # @example using a `Ronin::Arch` arch.
+  # @example using a `Ronin::Arch` arch (**deprecated**):
   #   0x41.pack(Arch.i686)
   #   # => "A\0\0\0"
   #
-  # @example specifying a custom address-length.
+  # @example specifying a custom address-length (**deprecated**):
   #   0x41.pack(Arch.ppc,2)
   #   # => "\0A"
   #
-  # @example using a `Array#pack` template String for the arch.
-  #   0x41.pack('L')
-  #   # => "A\0\0\0"
+  # @see http://rubydoc.info/stdlib/core/Array:pack
   #
-  # @see http://ruby-doc.org/core/classes/Array.html#M002222
-  #
   # @api public
   #
-  def pack(arch,address_length=nil)
-    if arch.kind_of?(String)
-      return [self].pack(arch)
-    end
+  def pack(*arguments)
+    argument = arguments.first
 
-    unless arch.respond_to?(:address_length)
-      raise(ArgumentError,"first argument to Ineger#pack must respond to address_length")
-    end
+    case argument
+    when String
+      [self].pack(argument)
+    when Symbol
+      unless Ronin::Binary::Template::INT_TYPES.include?(argument)
+        raise(ArgumentError,"unsupported integer type: #{argument}")
+      end
 
-    unless arch.respond_to?(:endian)
-      raise(ArgumentError,"first argument to Ineger#pack must respond to endian")
-    end
+      [self].pack(Ronin::Binary::Template::TYPES[argument])
+    else
+      # TODO: deprecate this calling convention
+      arch, address_length = arguments
 
-    address_length ||= arch.address_length
+      unless arch.respond_to?(:address_length)
+        raise(ArgumentError,"first argument to Ineger#pack must respond to address_length")
+      end
 
-    integer_bytes = bytes(address_length,arch.endian)
-    integer_bytes.map! { |b| b.chr }
+      unless arch.respond_to?(:endian)
+        raise(ArgumentError,"first argument to Ineger#pack must respond to endian")
+      end
 
-    return integer_bytes.join
+      address_length ||= arch.address_length
+
+      integer_bytes = bytes(address_length,arch.endian)
+      integer_bytes.map! { |b| b.chr }
+
+      return integer_bytes.join
+    end
   end
 
   #
   # @return [String]
   #   The hex escaped version of the Integer.