pwntools 0.1.0 → 1.2.1

data/lib/pwnlib/shellcraft/shellcraft.rb

@@ -0,0 +1,73 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'singleton'
+
+require 'pwnlib/context'
+require 'pwnlib/errors'
+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')].sort.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, **kwargs, &block)
+        mod = find_module_for(method)
+        return super if mod.nil?
+
+        mod.public_send(method, *args, **kwargs, &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
+          raise ::Pwnlib::Errors::UnsupportedArchError,
+                "Can't use shellcraft under architecture #{context.arch.inspect}."
+        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,67 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'time'
+
+require 'pwnlib/context'
+require 'pwnlib/errors'
+
+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 || ::Pwnlib::Context.context.timeout
+
+      @deadline = timeout == :forever ? :forever : Time.now + timeout
+
+      begin
+        yield
+      ensure
+        was_active = active?
+        @deadline = nil
+        raise ::Pwnlib::Errors::TimeoutError unless was_active
+      end
+    end
+  end
+end

data/lib/pwnlib/tubes/buffer.rb

@@ -0,0 +1,99 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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/process.rb

@@ -0,0 +1,155 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'pwnlib/errors'
+require 'pwnlib/tubes/tube'
+
+module Pwnlib
+  module Tubes
+    # Launch a process.
+    class Process < Tube
+      # Default options for {#initialize}.
+      DEFAULT_OPTIONS = {
+        env: ENV,
+        in: :pipe,
+        out: :pipe,
+        raw: true,
+        aslr: true
+      }.freeze
+
+      # Instantiate a {Pwnlib::Tubes::Process} object.
+      #
+      # @param [Array<String>, String] argv
+      #   List of arguments to pass to the spawned process.
+      #
+      # @option opts [Hash{String => String}] env (ENV)
+      #   Environment variables. By default, inherits from Ruby's environment.
+      # @option opts [Symbol] in (:pipe)
+      #   What kind of io should be used for +stdin+.
+      #   Candidates are: +:pipe+, +:pty+.
+      # @option opts [Symbol] out (:pipe)
+      #   What kind of io should be used for +stdout+.
+      #   Candidates are: +:pipe+, +:pty+.
+      #   See examples for more details.
+      # @option opts [Boolean] raw (true)
+      #   Set the created PTY to raw mode. i.e. disable echo and control characters.
+      #   If no pty is created, this has no effect.
+      # @option opts [Boolean] aslr (true)
+      #   If +false+ is given, the ASLR of the target process will be disabled via +setarch -R+.
+      # @option opts [Float?] timeout (nil)
+      #   See {Pwnlib::Tubes::Tube#initialize}.
+      #
+      # @example
+      #   io = Tubes::Process.new('ls')
+      #   io.gets
+      #   #=> "Gemfile\n"
+      #
+      #   io = Tubes::Process.new('ls', out: :pty)
+      #   io.gets
+      #   #=> "Gemfile       LICENSE\t\t\t   README.md  STYLE.md\t    git-hooks  pwntools.gemspec  test\n"
+      # @example
+      #   io = Tubes::Process.new('cat /proc/self/maps')
+      #   io.gets
+      #   #=> "55f8b8a10000-55f8b8a18000 r-xp 00000000 fd:00 9044035                    /bin/cat\n"
+      #   io.close
+      #
+      #   io = Tubes::Process.new('cat /proc/self/maps', aslr: false)
+      #   io.gets
+      #   #=> "555555554000-55555555c000 r-xp 00000000 fd:00 9044035                    /bin/cat\n"
+      #   io.close
+      # @example
+      #   io = Tubes::Process.new('env', env: { 'FOO' => 'BAR' })
+      #   io.gets
+      #   #=> "FOO=BAR\n"
+      def initialize(argv, **opts)
+        opts = DEFAULT_OPTIONS.merge(opts)
+        super(timeout: opts[:timeout])
+        argv = normalize_argv(argv, opts)
+        slave_i, slave_o = create_pipe(opts)
+        @pid = ::Process.spawn(opts[:env], *argv, in: slave_i, out: slave_o, unsetenv_others: true)
+        slave_i.close
+        slave_o.close unless slave_i == slave_o
+      end
+
+      # Close the IO.
+      #
+      # @param [:both, :recv, :read, :send, :write] direction
+      #   Disallow further read/write of the process.
+      #
+      # @return [void]
+      def shutdown(direction = :both)
+        close_io(normalize_direction(direction))
+      end
+
+      # Kill the process.
+      #
+      # @return [void]
+      def kill
+        shutdown
+        ::Process.kill('KILL', @pid)
+        ::Process.wait(@pid)
+      end
+      alias close kill
+
+      private
+
+      def io_out
+        @o
+      end
+
+      def close_io(dirs)
+        @o.close if dirs.include?(:read) && !@o.closed?
+        @i.close if dirs.include?(:write) && !@i.closed?
+      end
+
+      def normalize_argv(argv, opts)
+        # XXX(david942j): Set personality on child process will be better than using setarch
+        pre_cmd = opts[:aslr] ? '' : "setarch #{`uname -m`.strip} -R "
+        pre_cmd = pre_cmd.split if argv.is_a?(Array)
+        Array(pre_cmd + argv)
+      end
+
+      def create_pipe(opts)
+        if [opts[:in], opts[:out]].include?(:pty)
+          # Require only when we need it.
+          # This prevents broken on Windows, which has no pty support.
+          require 'io/console'
+          require 'pty'
+          mpty, spty = PTY.open
+          mpty.raw! if opts[:raw]
+        end
+        @o, slave_o = pipe(opts[:out], mpty, spty)
+        slave_i, @i = pipe(opts[:in], spty, mpty)
+        [slave_i, slave_o]
+      end
+
+      # @return [(IO, IO)]
+      #   IO pair.
+      def pipe(type, mst, slv)
+        case type
+        when :pipe then IO.pipe
+        when :pty then [mst, slv]
+        end
+      end
+
+      def send_raw(data)
+        @i.write(data)
+      rescue Errno::EIO, Errno::EPIPE, IOError
+        raise ::Pwnlib::Errors::EndOfTubeError
+      end
+
+      def recv_raw(size)
+        o, = IO.select([@o], [], [], @timeout)
+        return if o.nil?
+
+        @o.readpartial(size)
+      rescue Errno::EIO, Errno::EPIPE, IOError
+        raise ::Pwnlib::Errors::EndOfTubeError
+      end
+
+      def timeout_raw=(timeout)
+        @timeout = timeout
+      end
+    end
+  end
+end

data/lib/pwnlib/tubes/serialtube.rb

@@ -0,0 +1,114 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'rubyserial'
+
+require 'pwnlib/tubes/tube'
+
+module Pwnlib
+  module Tubes
+    # @!macro [new] raise_eof
+    #   @raise [Pwnlib::Errors::EndOfTubeError]
+    #     If the request is not satisfied when all data is received.
+
+    # Serial Connections
+    class SerialTube < Tube
+      # Instantiate a {Pwnlib::Tubes::SerialTube} object.
+      #
+      # @param [String] port
+      #   A device name for rubyserial to open, e.g. /dev/ttypUSB0
+      # @param [Integer] baudrate
+      #   Baud rate.
+      # @param [Boolean] convert_newlines
+      #   If +true+, convert any +context.newline+s to +"\\r\\n"+ before
+      #   sending to remote. Has no effect on bytes received.
+      # @param [Integer] bytesize
+      #   Serial character byte size. The '8' in '8N1'.
+      # @param [Symbol] parity
+      #   Serial character parity. The 'N' in '8N1'.
+      def initialize(port = nil, baudrate: 115_200,
+                     convert_newlines: true,
+                     bytesize: 8, parity: :none)
+        super()
+
+        # go hunting for a port
+        port ||= Dir.glob('/dev/tty.usbserial*').first
+        port ||= '/dev/ttyUSB0'
+
+        @convert_newlines = convert_newlines
+        @conn = Serial.new(port, baudrate, bytesize, parity)
+        @serial_timer = Timer.new
+      end
+
+      # Closes the active connection
+      def close
+        @conn.close if @conn && !@conn.closed?
+        @conn = nil
+      end
+
+      # Implementation of the methods required for tube
+      private
+
+      # Gets bytes over the serial connection until some bytes are received, or
+      # +@timeout+ has passed. It is an error for it to return no data in less
+      # than +@timeout+ seconds. It is ok for it to return some data in less
+      # time.
+      #
+      # @param [Integer] numbytes
+      #   An upper limit on the number of bytes to get.
+      #
+      # @return [String]
+      #   A string containing read bytes.
+      #
+      # @!macro raise_eof
+      def recv_raw(numbytes)
+        raise ::Pwnlib::Errors::EndOfTubeError if @conn.nil?
+
+        @serial_timer.countdown do
+          data = ''
+          begin
+            while @serial_timer.active?
+              data += @conn.read(numbytes - data.length)
+              break unless data.empty?
+
+              sleep 0.1
+            end
+            # XXX(JonathanBeverley): should we reverse @convert_newlines here?
+            return data
+          rescue RubySerial::Error
+            close
+            raise ::Pwnlib::Errors::EndOfTubeError
+          end
+        end
+      end
+
+      # Sends bytes over the serial connection. This call will block until all the bytes are sent or an error occurs.
+      #
+      # @param [String] data
+      #   A string of the bytes to send.
+      #
+      # @return [Integer]
+      #   The number of bytes successfully written.
+      #
+      # @!macro raise_eof
+      def send_raw(data)
+        raise ::Pwnlib::Errors::EndOfTubeError if @conn.nil?
+
+        data.gsub!(context.newline, "\r\n") if @convert_newlines
+        begin
+          @conn.write(data)
+        rescue RubySerial::Error
+          close
+          raise ::Pwnlib::Errors::EndOfTubeError
+        end
+      end
+
+      # Sets the +timeout+ to use for subsequent +recv_raw+ calls.
+      #
+      # @param [Float] timeout
+      def timeout_raw=(timeout)
+        @serial_timer.timeout = timeout
+      end
+    end
+  end
+end