Feature/serialtube

Gemfile.lock

@@ -9,6 +9,7 @@ PATH
       method_source (~> 0.9)
       rainbow (>= 2.2, < 4.0)
       ruby2ruby (~> 2.4)
+      rubyserial (~> 0.5)
 
 GEM
   remote: https://www.rubygems.org/
@@ -50,6 +51,8 @@ GEM
       sexp_processor (~> 4.6)
     ruby_parser (3.11.0)
       sexp_processor (~> 4.9)
+    rubyserial (0.5.0)
+      ffi (~> 1.9, >= 1.9.3)
     sexp_processor (4.10.1)
     simplecov (0.16.1)
       docile (~> 1.1)

lib/pwnlib/pwn.rb

@@ -12,6 +12,7 @@
 require 'pwnlib/logger'
 require 'pwnlib/reg_sort'
 require 'pwnlib/shellcraft/shellcraft'
+require 'pwnlib/tubes/serialtube'
 require 'pwnlib/tubes/sock'
 
 require 'pwnlib/util/cyclic'

lib/pwnlib/timer.rb

@@ -54,8 +54,9 @@ def countdown(timeout = nil)
       begin
         yield
       ensure
-        raise ::Pwnlib::Errors::TimeoutError unless active?
+        was_active = active?
         @deadline = nil
+        raise ::Pwnlib::Errors::TimeoutError unless was_active
       end
     end
   end

lib/pwnlib/tubes/serialtube.rb

@@ -0,0 +1,107 @@
+# encoding: ASCII-8BIT
+
+require 'rubyserial'
+
+require 'pwnlib/tubes/tube'
+
+module Pwnlib
+  module Tubes
+    # 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+, will convert local +context.newline+ to +"\\r\\n"+ for remote
+      # @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 = 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 EOFError 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
+            # TODO: should we reverse @convert_newlines here?
+            return data
+          rescue RubySerial::Error
+            close
+            raise EOFError
+          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 EOFError if @conn.nil?
+
+        data.gsub!(context.newline, "\r\n") if @convert_newlines
+        begin
+          return @conn.write(data)
+        rescue RubySerial::Error
+          close
+          raise EOFError
+        end
+      end
+
+      # Sets the +timeout+ to use for subsequent +recv_raw+ calls.
+      #
+      # @param [Integer] timeout
+      def timeout_raw=(timeout)
+        @serial_timer.timeout = timeout
+      end
+    end
+  end
+end

lib/pwnlib/tubes/tube.rb

@@ -51,6 +51,9 @@ def initialize(timeout: nil)
       # @return [String]
       #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
       #   waiting.
+      #
+      # @!macro raise_eof
+      # @!macro raise_timeout
       def recv(num_bytes = nil, timeout: nil)
         return '' if @buffer.empty? && !fillbuffer(timeout: timeout)
         @buffer.get(num_bytes)
@@ -101,13 +104,17 @@ def recvpred(timeout: nil)
 
               c = recv(1)
 
-              return '' if c.empty?
+              next if c.empty?
               data << c
             end
             data.slice!(0..-1)
-          ensure
+          # rubocop:disable Style/RescueStandardError
+          # Justification: go home rubocop, you're drunk.
+          rescue
             unrecv(data)
+            raise
           end
+          # rubocop:enable Style/RescueStandardError
         end
       end
 
@@ -159,7 +166,7 @@ def recvuntil(delims, drop: false, timeout: nil)
             while @timer.active?
               s = recv(1)
 
-              return '' if s.empty?
+              next if s.empty?
               matching << s
 
               sidx = matching.size

pwntools.gemspec

@@ -33,6 +33,7 @@ Gem::Specification.new do |s|
   s.add_runtime_dependency 'method_source', '~> 0.9'
   s.add_runtime_dependency 'rainbow', '>= 2.2', '< 4.0'
   s.add_runtime_dependency 'ruby2ruby', '~> 2.4'
+  s.add_runtime_dependency 'rubyserial', '~> 0.5'
 
   # TODO(david942j): check why ruby crash during testing if upgrade minitest to 5.10.2/3
   s.add_development_dependency 'minitest', '= 5.10.1'

test/tubes/serialtube_test.rb

@@ -0,0 +1,139 @@
+# encoding: ASCII-8BIT
+
+require 'open3'
+
+require 'test_helper'
+
+require 'pwnlib/tubes/serialtube'
+
+class SerialTest < MiniTest::Test
+  include ::Pwnlib::Tubes
+
+  def skip_windows
+    skip 'Not test tube/serialtube on Windows' if TTY::Platform.new.windows?
+  end
+
+  def open_pair
+    Open3.popen3('socat -d -d pty,raw,echo=0 pty,raw,echo=0') do |_i, _o, stderr, thread|
+      devs = []
+      2.times do
+        devs << stderr.readline.chomp.split.last
+        # First pattern matches Linux, second is macOS
+        raise IOError, 'Could not create serial crosslink' if devs.last !~ %r{^(/dev/pts/[0-9]+|/dev/ttys[0-9]+)$}
+      end
+
+      serial = SerialTube.new devs[1], convert_newlines: false
+
+      begin
+        File.open devs[0], 'r+' do |file|
+          file.set_encoding 'default'.encoding
+          yield file, serial
+        end
+      ensure
+        Process.kill('SIGTERM', thread.pid)
+      end
+    end
+  end
+
+  def random_string(length)
+    Random.rand(36**length).to_s(36).encode('default'.encoding)
+  end
+
+  def test_recv
+    skip_windows
+    open_pair do |file, serial|
+      # recv, recvline
+      rs = random_string 24
+      file.puts rs
+      result = serial.recv 8, timeout: 1
+
+      assert_equal(rs[0...8], result)
+      result = serial.recv 8
+      assert_equal(rs[8...16], result)
+      result = serial.recvline.chomp
+      assert_equal(rs[16..-1], result)
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvpred
+      rs = random_string 12
+      file.print rs
+      result = serial.recvpred do |data|
+        data[-6..-1] == rs[-6..-1]
+      end
+      assert_equal rs, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvn
+      rs = random_string 6
+      file.print rs
+      result = ''
+      assert_raises(Pwnlib::Errors::TimeoutError) do
+        result = serial.recvn 120, timeout: 1
+      end
+      assert_empty result
+      file.print rs
+      result = serial.recvn 12
+      assert_equal rs * 2, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvuntil
+      rs = random_string 12
+      file.print rs + '|'
+      result = serial.recvuntil('|').chomp('|')
+      assert_equal rs, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # gets
+      rs = random_string 24
+      file.puts rs
+      result = serial.gets 12
+      assert_equal rs[0...12], result
+      result = serial.gets.chomp
+      assert_equal rs[12..-1], result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+    end
+  end
+
+  def test_send
+    skip_windows
+    open_pair do |file, serial|
+      # send, sendline
+      rs = random_string 24
+      # rubocop:disable Style/Send
+      # Justification: This isn't Object#send, false positive.
+      serial.send rs[0...12]
+      # rubocop:enable Style/Send
+      serial.sendline rs[12...24]
+      result = file.readline.chomp
+      assert_equal rs, result
+
+      # puts
+      r1 = random_string 4
+      r2 = random_string 4
+      r3 = random_string 4
+      serial.puts r1, r2, r3
+      result = ''
+      3.times do
+        result += file.readline.chomp
+      end
+      assert_equal r1 + r2 + r3, result
+    end
+  end
+
+  def test_close
+    skip_windows
+    open_pair do |_file, serial|
+      serial.close
+      assert_raises(EOFError) { serial.puts(514) }
+      assert_raises(EOFError) { serial.puts(514) }
+      assert_raises(EOFError) { serial.recv }
+      assert_raises(EOFError) { serial.recv }
+      assert_raises(ArgumentError) { serial.close(:hh) }
+    end
+  end
+end

travis/install.sh

@@ -49,6 +49,9 @@ setup_osx()
   install_keystone_from_source
   ln -s keystone/build/llvm/lib/libkeystone.dylib libkeystone.dylib # hack, don't know why next line has no effect
   # export DYLD_LIBRARY_PATH=$TRAVIS_BUILD_DIR/keystone/build/llvm/lib:$DYLD_LIBRARY_PATH
+
+  # install socat
+  brew install socat
 }
 
 if [[ "$TRAVIS_OS_NAME" == "osx" ]]; then