pwntools 0.1.0 → 1.0.0

data/lib/pwnlib/shellcraft/generators/x86/linux/syscall.rb

@@ -0,0 +1,52 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/shellcraft/generators/x86/common/setregs'
+require 'pwnlib/shellcraft/generators/x86/linux/linux'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module X86
+        module Linux
+          # Assembly of +syscall+.
+          #
+          # @example
+          #   context.arch = 'i386'
+          #   puts shellcraft.syscall('SYS_open', 'esp', 0, 0)
+          #   # /* call open("esp", 0, 0) */
+          #   # push 5 /* (SYS_open) */
+          #   # pop eax
+          #   # mov ebx, esp
+          #   # xor ecx, ecx /* 0 */
+          #   # cdq /* edx=0 */
+          #   # int 0x80
+          #   #=> nil
+          def syscall(*arguments)
+            abi = ::Pwnlib::ABI::ABI.syscall
+            syscall, arg0, arg1, arg2, arg3, arg4, arg5 = arguments
+            if syscall.respond_to?(:to_s) && syscall.to_s.start_with?('SYS_')
+              syscall_repr = syscall.to_s[4..-1] + '(%s)'
+              args = []
+            else
+              syscall_repr = 'syscall(%s)'
+              args = [syscall ? syscall.inspect : '?']
+            end
+            # arg0 to arg5
+            1.upto(6) do |i|
+              args.push(arguments[i] ? arguments[i].inspect : '?')
+            end
+
+            args.pop while args.last == '?'
+            syscall_repr = format(syscall_repr, args.join(', '))
+            registers = abi.register_arguments
+            arguments = [syscall, arg0, arg1, arg2, arg3, arg4, arg5]
+            reg_ctx = registers.zip(arguments).to_h
+            cat "/* call #{syscall_repr} */"
+            cat Common.setregs(reg_ctx) if arguments.any? { |v| !v.nil? }
+            cat abi.syscall_str
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pwnlib/shellcraft/registers.rb

@@ -0,0 +1,145 @@
+require 'pwnlib/context'
+
+module Pwnlib
+  module Shellcraft
+    # Define register names and methods for shellcode generators.
+    module Registers
+      X86_BASEREGS = %w(ax cx dx bx sp bp si di ip).freeze
+
+      I386 = (X86_BASEREGS.map { |r| "e#{r}" } +
+             X86_BASEREGS +
+             %w(eflags cs ss ds es fs gs)).freeze
+
+      AMD64 = (X86_BASEREGS.map { |r| "r#{r}" } +
+               (8..15).map { |r| "r#{r}" } +
+               (8..15).map { |r| "r#{r}d" } +
+               I386).freeze
+
+      # x86 registers in decreasing size
+      X86_ORDERED = ([
+        %w(rax eax ax al),
+        %w(rbx ebx bx bl),
+        %w(rcx ecx cx cl),
+        %w(rdx edx dx dl),
+        %w(rdi edi di),
+        %w(rsi esi si),
+        %w(rbp ebp bp),
+        %w(rsp esp sp)
+      ] + (8..15).map { |r| ['', 'd', 'w', 'b'].map { |t| "r#{r}#{t}" } }).freeze
+
+      # class Register, currently only supports i386 and amd64.
+      class Register
+        # @return [String]
+        #   Register's name.
+        attr_reader :name
+        attr_reader :bigger, :smaller, :ff00, :is64bit, :native64, :native32, :xor
+        attr_reader :size, :sizes
+
+        # Instantiate a {Register} object.
+        #
+        # Create a register by its name and size (in bits) for fetching other information. For example, for register
+        # 'ax', +#bigger+ contains 'rax' and 'eax'.
+        #
+        # Normally you don't need to create any {Register} object, use {.get_register} to get register by name.
+        #
+        # @param [String] name
+        #   Register's name.
+        # @param [Integer] size
+        #   Register size in bits.
+        #
+        # @example
+        #   Register.new('rax', 64)
+        #   Register.new('bx', 16)
+        def initialize(name, size)
+          @name = name
+          @size = size
+          X86_ORDERED.each do |row|
+            next unless row.include?(name)
+            @bigger = row[0, row.index(name)]
+            @smaller = row[(row.index(name) + 1)..-1]
+            @native64 = row[0]
+            @native32 = row[1]
+            @sizes = row.each_with_object({}).with_index { |(r, h), i| h[64 >> i] = r }
+            @xor = @sizes[[size, 32].min]
+            break
+          end
+          @ff00 = name[1] + 'h' if @size >= 32 && @name.end_with?('x')
+          @is64bit = true if @name.start_with?('r')
+        end
+
+        def bits
+          size
+        end
+
+        def bytes
+          bits / 8
+        end
+
+        def fits(value)
+          size >= Registers.bits_required(value)
+        end
+
+        def to_s
+          name
+        end
+
+        def inspect
+          format('Register(%s)', name)
+        end
+      end
+
+      module_function
+
+      def registers
+        {
+          [32, 'i386', 'linux'] => ::Pwnlib::Shellcraft::Registers::I386,
+          [64, 'amd64', 'linux'] => ::Pwnlib::Shellcraft::Registers::AMD64
+        }[[context.bits, context.arch, context.os]]
+      end
+
+      INTEL = (X86_ORDERED.each_with_object({}) do |row, obj|
+        row.each_with_index do |reg, i|
+          obj[reg] = Register.new(reg, 64 >> i)
+        end
+      end).freeze
+
+      # Get a {Register} object by name.
+      #
+      # @param [String, Symbol, Register] name
+      #   The name of register.
+      #   If +name+ is already a {Register} object, +name+ itself will be returned.
+      #
+      # @return [Register?]
+      #   Get the register with name +name+.
+      #
+      # @example
+      #   Registers.get_register('eax')
+      #   #=> Register(eax)
+      #   Registers.get_register(:ebx)
+      #   #=> Register(ebx)
+      #   Registers.get_register('xdd')
+      #   #=> nil
+      def get_register(name)
+        return name if name.instance_of?(Register)
+        return INTEL[name.to_s] if name.instance_of?(String) || name.instance_of?(Symbol)
+        nil
+      end
+
+      def register?(obj)
+        !get_register(obj).nil?
+      end
+
+      def bits_required(value)
+        bits = 0
+        value = value.abs
+        while value > 0
+          bits += 8
+          value >>= 8
+        end
+        bits
+      end
+
+      include ::Pwnlib::Context
+    end
+  end
+end

data/lib/pwnlib/shellcraft/shellcraft.rb

@@ -0,0 +1,67 @@
+# encoding: ASCII-8BIT
+
+require 'singleton'
+
+require 'pwnlib/context'
+require 'pwnlib/logger'
+
+module Pwnlib
+  # Implement shellcraft!
+  #
+  # All shellcode generators are defined under generators/*.
+  # While typing +Shellcraft::Generators::I386::Linux.sh+ is too annoying, we define an instance +shellcraft+ in this
+  # module, which let user invoke +shellcraft.sh+ directly.
+  module Shellcraft
+    # Singleton class.
+    class Shellcraft
+      include ::Singleton
+
+      # All files under generators/ will be required.
+      def initialize
+        Dir[File.join(__dir__, 'generators', '**', '*.rb')].each do |f|
+          require f
+        end
+      end
+
+      # Will search modules/methods under {Shellcraft::Generators} according to current arch and os.
+      # i.e. +Shellcraft::Generators::${arch}::<Common|${os}>.${method}+.
+      #
+      # With this method, +context.local(arch: 'amd64') { shellcraft.sh }+ will invoke
+      # {Shellcraft::Generators::Amd64::Linux#sh}.
+      def method_missing(method, *args, &block)
+        mod = find_module_for(method)
+        return super if mod.nil?
+        mod.public_send(method, *args, &block)
+      end
+
+      # For +respond_to?+.
+      def respond_to_missing?(method, include_private = false)
+        return true if find_module_for(method)
+        super
+      end
+
+      private
+
+      # @return [Module?]
+      #   +nil+ for not found.
+      def find_module_for(method)
+        begin
+          arch_module = ::Pwnlib::Shellcraft::Generators.const_get(context.arch.capitalize)
+        rescue NameError
+          log.error("Can't use shellcraft under architecture #{context.arch.inspect}.")
+          return nil
+        end
+        # try search in Common module
+        common_module = arch_module.const_get(:Common)
+        return common_module if common_module.singleton_methods.include?(method)
+        # search in ${os} module
+        os_module = arch_module.const_get(context.os.capitalize)
+        return os_module if os_module.singleton_methods.include?(method)
+        nil
+      end
+
+      include ::Pwnlib::Context
+      include ::Pwnlib::Logger
+    end
+  end
+end

data/lib/pwnlib/timer.rb

@@ -0,0 +1,60 @@
+# encoding: ASCII-8BIT
+
+require 'time'
+
+require 'pwnlib/context'
+
+module Pwnlib
+  # A simple timer class.
+  # TODO(Darkpi): Python pwntools seems to have many unreasonable codes in this class,
+  #               not sure of the use case of this, check if everything is coded as
+  #               intended after we have some use cases. (e.g. sock)
+  # NOTE(Darkpi): This class is actually quite weird, and expected to be used only in tubes.
+  class Timer
+    # @diff We just use nil for default and :forever for forever.
+
+    def initialize(timeout = nil)
+      @deadline = nil
+      @timeout = timeout
+    end
+
+    def started?
+      @deadline
+    end
+
+    def active?
+      started? && (@deadline == :forever || Time.now < @deadline)
+    end
+
+    def timeout
+      return @timeout || Pwnlib::Context.context.timeout unless started?
+      @deadline == :forever ? :forever : [@deadline - Time.now, 0].max
+    end
+
+    def timeout=(timeout)
+      raise "Can't change timeout when countdown" if started?
+      @timeout = timeout
+    end
+
+    # @diff We do NOT allow nested countdown with non-default value. This simplifies thing a lot.
+    # NOTE(Darkpi): timeout = nil means default value for the first time, and nop after that.
+    def countdown(timeout = nil)
+      raise ArgumentError, 'Need a block for countdown' unless block_given?
+      if started?
+        return yield if timeout.nil?
+        raise 'Nested countdown not permitted'
+      end
+
+      timeout ||= @timeout
+      timeout ||= Pwnlib::Context.context.timeout
+
+      @deadline = timeout == :forever ? :forever : Time.now + timeout
+
+      begin
+        yield
+      ensure
+        @deadline = nil
+      end
+    end
+  end
+end

data/lib/pwnlib/tubes/buffer.rb

@@ -0,0 +1,96 @@
+# encoding: ASCII-8BIT
+
+module Pwnlib
+  module Tubes
+    # Buffer that support deque-like string operations.
+    class Buffer
+      # Instantiate a {Pwnlib::Tubes::Buffer} object.
+      def initialize
+        @data = []
+        @size = 0
+      end
+
+      attr_reader :size
+      alias length size
+
+      # Check whether the buffer is empty.
+      #
+      # @return [Boalean]
+      #   Returns true if contains no elements.
+      def empty?
+        size.zero?
+      end
+
+      # Python __contains__ and index is only correct with single-char input, which doesn't seems to
+      # be useful, and they're not used anywhere either. Just ignore them for now.
+
+      # Adds data to the buffer.
+      #
+      # @param [String] data
+      #   Data to add.
+      def add(data)
+        case data
+        when Buffer
+          @data.concat(data.data)
+        else
+          data = data.to_s.dup
+          return if data.empty?
+          @data << data
+        end
+        @size += data.size
+        self
+      end
+      alias << add
+
+      # Places data at the front of the buffer.
+      #
+      # @param [String] data
+      #   Data to place at the beginning of the buffer.
+      def unget(data)
+        case data
+        when Buffer
+          @data.unshift(*data.data)
+        else
+          data = data.to_s.dup
+          return if data.empty?
+          @data.unshift(data)
+        end
+        @size += data.size
+        self
+      end
+
+      # Retrieves bytes from the buffer.
+      #
+      # @param [Integer] n
+      #   Maximum number of bytes to fetch.
+      #
+      # @return [String]
+      #   Data as string.
+      def get(n = nil)
+        if n.nil? || n >= size
+          data = @data.join
+          @size = 0
+          @data = []
+        else
+          have = 0
+          idx = 0
+          while have < n
+            have += @data[idx].size
+            idx += 1
+          end
+          data = @data.slice!(0...idx).join
+          if have > n
+            extra = data.slice!(n..-1)
+            @data.unshift(extra)
+          end
+          @size -= n
+        end
+        data
+      end
+
+      protected
+
+      attr_reader :data
+    end
+  end
+end

data/lib/pwnlib/tubes/sock.rb

@@ -0,0 +1,95 @@
+# encoding: ASCII-8BIT
+
+require 'socket'
+
+require 'pwnlib/tubes/tube'
+
+module Pwnlib
+  module Tubes
+    # Socket!
+    class Sock < Tube
+      # Instantiate a {Pwnlib::Tubes::Sock} object.
+      #
+      # @param [String] host
+      #   The host to connect.
+      # @param [Integer] port
+      #   The port to connect.
+      def initialize(host, port)
+        super()
+        @sock = TCPSocket.new(host, port)
+        @sock.binmode
+        @timeout = nil
+        @closed = { recv: false, send: false }
+      end
+
+      def io
+        @sock
+      end
+      alias sock io
+
+      # Close the TCPSocket if no arguments passed.
+      # Or close the direction in +sock+.
+      #
+      # @param [:both, :recv, :read, :send, :write] direction
+      #   * Close the TCPSocket if +:both+ or no arguments passed.
+      #   * Disallow further read in +sock+ if +:recv+ or +:read+ passed.
+      #   * Disallow further write in +sock+ if +:send+ or +:write+ passed.
+      #
+      # @diff In pwntools-python, method +shutdown(direction)+ is for closing socket one side,
+      #   +close()+ is for closing both side. We merge these two methods into one here.
+      def close(direction = :both)
+        case direction
+        when :both
+          return if @sock.closed?
+          @closed[:recv] = @closed[:send] = true
+          @sock.close
+        when :recv, :read
+          shutdown(:recv)
+        when :send, :write
+          shutdown(:send)
+        else
+          raise ArgumentError, 'Only allow :both, :recv, :read, :send and :write passed'
+        end
+      end
+
+      private
+
+      def shutdown(direction)
+        return if @closed[direction]
+        @closed[direction] = true
+
+        if direction.equal?(:recv)
+          @sock.close_read
+        elsif direction.equal?(:send)
+          @sock.close_write
+        end
+      end
+
+      def timeout_raw=(timeout)
+        @timeout = timeout == :forever ? nil : timeout
+      end
+
+      def send_raw(data)
+        raise EOFError if @closed[:send]
+        begin
+          @sock.write(data)
+        rescue Errno::EPIPE, Errno::ECONNRESET, Errno::ECONNREFUSED
+          shutdown(:send)
+          raise EOFError
+        end
+      end
+
+      def recv_raw(size)
+        raise EOFError if @closed[:recv]
+        begin
+          rs, = IO.select([@sock], [], [], @timeout)
+          return if rs.nil?
+          return @sock.readpartial(size)
+        rescue Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::ECONNABORTED, EOFError
+          shutdown(:recv)
+          raise EOFError
+        end
+      end
+    end
+  end
+end