refactor builds' format

.rubocop.yml

@@ -13,3 +13,6 @@ Rails:
 Metrics/LineLength:
   Enabled: true
   Max: 120
+
+Style/FileName:
+  Enabled: false

lib/one_gadget.rb

@@ -3,12 +3,16 @@
 # @author david942j
 module OneGadget
   class << self
-    # The man entry of gem +one_gadget+
-    # @option [String] file
+    # The man entry of gem +one_gadget+.
+    # If want to find gadgets from file, it will search gadgets by its
+    # build id first.
+    #
+    # @param [String] filepath
     #   The relative path of +libc.so.6+.
-    # @option [String] build_id
+    # @param [String] build_id
     #   The BuildID of target +libc.so.6+.
-    # @return [Array]
+    # @return [Array<OneGadget::Gadget::Gadget>, Array<Integer>]
+    #   The gadgets found.
     # @example
     #   OneGadget.gadgets(filepath: './libc.so.6')
     #   OneGadget.gadgets(build_id: '60131540dadc6796cab33388349e6e4e68692053')

lib/one_gadget/abi.rb

@@ -5,6 +5,7 @@ module OneGadget
     module ClassMethods
       LINUX_X86_32 = %w(eax ebx ecx edx edi esi ebp esp).freeze
       LINUX_X86_64 = LINUX_X86_32 + %w(rax rbx rcx rdx rdi rsi rbp rsp) + 7.upto(15).map { |i| "r#{i}" }
+      # Only support x86-64 now.
       def registers
         LINUX_X86_64
       end

lib/one_gadget/builds/60131540dadc6796cab33388349e6e4e68692053.rb

@@ -1,13 +0,0 @@
-require 'one_gadget/gadget'
-# Ubuntu GLIBC 2.23-0ubuntu5
-# ELF 64-bit LSB shared object, x86-64
-build_id = File.basename(__FILE__, '.rb')
-OneGadget::Gadget.define(build_id) do |g|
-  g.offset = 0x4526a
-  g.constraints = ['[rsp+0x30]=NULL']
-end
-
-OneGadget::Gadget.define(build_id) do |g|
-  g.offset = 0xef6c4
-  g.constraints = ['[rsp+0x50]=NULL']
-end

lib/one_gadget/builds/libc-2.23-60131540dadc6796cab33388349e6e4e68692053.rb

@@ -0,0 +1,6 @@
+require 'one_gadget/gadget'
+# Ubuntu GLIBC 2.23-0ubuntu5
+# ELF 64-bit LSB shared object, x86-64
+build_id = File.basename(__FILE__, '.rb').split('-').last
+OneGadget::Gadget.add(build_id, 0x4526a, constraints: ['[rsp+0x30] == NULL'])
+OneGadget::Gadget.add(build_id, 0xef6c4, constraints: ['[rsp+0x50] == NULL'])

lib/one_gadget/fetcher.rb

@@ -6,10 +6,14 @@ module OneGadget
   module Fetcher
     # Define class methods here.
     module ClassMethods
+      # Fetch one-gadget offsets of this build id.
       # @param [String] build_id The targets' BuildID.
-      # @option [Boolean] details
+      # @param [Boolean] details
       #   If needs to return the gadgets' constraints.
-      # @return [Array]
+      # @return [Array<Integer>, Array<OneGadget::Gadget::Gadget>]
+      #   If +details+ is +false+, +Array<Integer>+ is returned, which
+      #   contains offset only.
+      #   Otherwise, array of gadgets is returned.
       def from_build_id(build_id, details: false)
         if (build_id =~ /\A#{OneGadget::Helper::BUILD_ID_FORMAT}\Z/).nil?
           raise ArgumentError, format('invalid BuildID format: %p', build_id)

lib/one_gadget/gadget.rb

@@ -5,13 +5,23 @@ module OneGadget
   module Gadget
     # Information of a gadget.
     class Gadget
+      # @return [Integer] The gadget's address offset.
       attr_accessor :offset
+      # @return [Array<String>] The constraints need for this gadget.
       attr_accessor :constraints
-      def constraints=(value)
-        raise ArgumentError, value unless value.is_a?(Array)
-        @constraints = value
+
+      # Initialize method of {Gadget} instance.
+      # @param [Integer] offset The relative address offset of this gadget.
+      # @option options [Array<String>] :constraints
+      #   The constraints need for this gadget. Defaults to +[]+.
+      # @example
+      #   OneGadget::Gadget::Gadget.new(0x12345, constraints: ['rax == 0'])
+      def initialize(offset, **options)
+        @offset = offset
+        @constraints = options[:constraints] || []
       end
 
+      # Show gadget in a pretty way.
       def inspect
         str = format("#{OneGadget::Helper.colorize('offset', sev: :sym)}: 0x%x\n", offset)
         unless constraints.nil?
@@ -22,29 +32,37 @@ module OneGadget
         OneGadget::ABI.registers.each { |reg| str.gsub!(reg, OneGadget::Helper.colorize(reg, sev: :reg)) }
         str + "\n"
       end
-
-      # Only check if +offset+ being set now.
-      def self_check!
-        raise ArgumentError, format('invalid offset: %p', offset) unless offset.is_a?(Integer) && offset >= 0
-      end
     end
 
     # Define class methods here.
     module ClassMethods
       BUILDS_PATH = File.join(File.dirname(__FILE__), 'builds').freeze
-      CACHE = Hash.new { |h, k| h[k] = [] }
+      BUILDS = Hash.new { |h, k| h[k] = [] }
+      # Get gadgets from pre-defined corpus.
+      # @param [String] build_id Desired build id.
+      # @return [Array<Gadget::Gadget>] Gadgets.
       def builds(build_id)
-        return CACHE[build_id] if CACHE.key?(build_id)
-        return [] unless File.exist?(File.join(BUILDS_PATH, build_id + '.rb'))
-        require File.join(BUILDS_PATH, build_id + '.rb')
-        CACHE[build_id]
+        require_all if BUILDS.empty?
+        return BUILDS[build_id] if BUILDS.key?(build_id)
+        # TODO: fetch remote builds information.
+        []
       end
 
-      def define(build_id)
-        g = OneGadget::Gadget::Gadget.new
-        yield(g)
-        g.self_check!
-        CACHE[build_id] << g
+      # Add a gadget, for scripts in builds/ to use.
+      # @param [String] build_id The target's build id.
+      # @param [Integer] offset The relative address offset of this gadget.
+      # @param [Hash] options See {Gadget::Gadget#initialize} for more information.
+      # @return [void]
+      def add(build_id, offset, **options)
+        BUILDS[build_id] << OneGadget::Gadget::Gadget.new(offset, **options)
+      end
+
+      private
+
+      def require_all
+        Dir.glob(File.join(BUILDS_PATH, '**', '*.rb')).each do |dic|
+          require dic
+        end
       end
     end
     extend ClassMethods

lib/one_gadget/helper.rb

@@ -29,10 +29,12 @@ module OneGadget
         bid.first
       end
 
+      # Disable colorize
       def color_off!
         @disable_color = true
       end
 
+      # Enable colorize
       def color_on!
         @disable_color = false
       end
@@ -47,7 +49,7 @@ module OneGadget
         sym: "\e[38;5;229m", # pry like
       }.freeze
 
-      # Wrapper color codes for for pretty inspect.
+      # Wrapper color codes for pretty inspect.
       # @param [String] str Contents to colorize.
       # @option [Symbol] sev Specific which kind of color want to use, valid symbols are defined in +COLOR_CODE+.
       # @return [String] Wrapper with color codes.

spec/gadget_spec.rb

@@ -4,18 +4,15 @@ describe OneGadget::Gadget do
   before(:all) do
     @build_id = 'fake_id'
     OneGadget::Helper.color_off! # disable colorize for easy testing.
-    OneGadget::Gadget.define(@build_id) do |g|
-      g.offset = 0x1234
-      g.constraints = ['[rsp+0x30]=NULL', 'rax=0']
-    end
+    OneGadget::Gadget.add(@build_id, 0x1234, constraints: ['[rsp+0x30] == NULL', 'rax == 0'])
   end
 
   it 'inspect' do
     expect(OneGadget::Gadget.builds(@build_id).map(&:inspect).join).to eq <<-EOS
 offset: 0x1234
 constraints:
-  [rsp+0x30]=NULL
-  rax=0
+  [rsp+0x30] == NULL
+  rax == 0
     EOS
   end
 end

spec/one_gadget_spec.rb

@@ -2,6 +2,8 @@ require 'one_gadget'
 
 describe 'one_gadget' do
   before(:all) do
+    # To force require all again.
+    OneGadget::Gadget::ClassMethods::BUILDS.clear
     @build_id = '60131540dadc6796cab33388349e6e4e68692053'
     @libcpath = File.join(File.dirname(__FILE__), 'data', 'libc-2.19.so')
   end
@@ -11,12 +13,12 @@ describe 'one_gadget' do
   end
 
   describe 'from build id' do
-    it 'from build id' do
+    it 'normal' do
       # only check not empty because the gadgets might add frequently.
       expect(OneGadget.gadgets(build_id: @build_id)).not_to be_empty
     end
 
-    it 'invalid id' do
+    it 'invalid' do
       expect { OneGadget.gadgets(build_id: '^_^') }.to raise_error(ArgumentError, 'invalid BuildID format: "^_^"')
     end
   end