lib/pwnlib/tubes/tube.rb
# encoding: ASCII-8BIT
# frozen_string_literal: true
require 'pwnlib/context'
require 'pwnlib/errors'
require 'pwnlib/logger'
require 'pwnlib/timer'
require 'pwnlib/tubes/buffer'
require 'pwnlib/util/hexdump'
module Pwnlib
# Module that contains all kinds of tubes.
module Tubes
# @!macro [new] drop_definition
# @param [Boolean] drop
# Whether drop the ending.
#
# @!macro [new] timeout_definition
# @param [Float] timeout
# Any positive floating number, indicates timeout in seconds.
# Using +context.timeout+ if +timeout+ equals to +nil+.
#
# @!macro [new] send_return_definition
# @return [Integer]
# Returns the number of bytes had been sent.
#
# @!macro [new] raise_eof
# @raise [Pwnlib::Errors::EndOfTubeError]
# If the request is not satisfied when all data is received.
#
# @!macro [new] raise_timeout
# @raise [Pwnlib::Errors::TimeoutError]
# If the request is not satisfied when timeout exceeded.
# Things common to all tubes (sockets, tty, ...)
class Tube
# Receive 4096 bytes each time.
BUFSIZE = 4096
# Instantiate a {Pwnlib::Tubes::Tube} object.
#
# @!macro timeout_definition
def initialize(timeout: nil)
@timer = Timer.new(timeout)
@buffer = Buffer.new
end
# Receives up to +num_bytes+ bytes of data from the tube, and returns as soon as any quantity
# of data is available.
#
# @param [Integer] num_bytes
# The maximum number of bytes to receive.
# @!macro timeout_definition
#
# @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)
end
alias read recv
# Puts the specified data back at the beginning of the receive buffer.
#
# @param [String] data
# A string to put back.
#
# @return [Integer]
# The length of the put back data.
def unrecv(data)
@buffer.unget(data)
data.size
end
# Receives one byte at a time from the tube, until the predicate evaluates to +true+.
#
# @!macro timeout_definition
#
# @return [String]
# A string contains bytes received from the tube, or +''+ if a timeout occurred while
# waiting.
#
# @yield
# A predicate to evaluate whether the data satisfy the condition.
#
# @yieldparam [String] data
# A string data to be validated by the predicate.
#
# @yieldreturn [Boolean]
# Whether the data satisfy the condition.
#
# @raise [ArgumentError]
# If the block is not given.
#
# @!macro raise_eof
# @!macro raise_timeout
def recvpred(timeout: nil)
raise ArgumentError, 'Need a block for recvpred' unless block_given?
@timer.countdown(timeout) do
data = +''
begin
until yield(data)
return '' unless @timer.active?
c = recv(1)
return '' if c.empty?
data << c
end
data.slice!(0..-1)
ensure
unrecv(data)
end
end
end
# Receives exactly +num_bytes+ bytes.
# If the request is not satisfied before +timeout+ seconds pass, all data is buffered and an
# empty string +''+ is returned.
#
# @param [Integer] num_bytes
# The number of bytes to receive.
# @!macro timeout_definition
#
# @return [String]
# A string contains bytes received from the tube, or +''+ if a timeout occurred while
# waiting.
#
# @!macro raise_eof
# @!macro raise_timeout
def recvn(num_bytes, timeout: nil)
@timer.countdown(timeout) do
fillbuffer while @timer.active? && @buffer.size < num_bytes
@buffer.size >= num_bytes ? @buffer.get(num_bytes) : ''
end
end
alias readn recvn
# Receives data until one of +delims+ is encountered. If the request is not satisfied before
# +timeout+ seconds pass, all data is buffered and an empty string is returned.
#
# @param [Array<String>] delims
# String of delimiters characters, or list of delimiter strings.
# @!macro drop_definition
# @!macro timeout_definition
#
# @return [String]
# A string contains bytes, which ends string in +delims+, received from the tube.
#
# @!macro raise_eof
# @!macro raise_timeout
#
# @diff We return the string that ends the earliest, rather then starts the earliest,
# since the latter can't be done greedly. Still, it would be bad to call this
# for case with ambiguity.
def recvuntil(delims, drop: false, timeout: nil)
delims = Array(delims)
max_len = delims.map(&:size).max
@timer.countdown(timeout) do
data = Buffer.new
matching = +''
begin
while @timer.active?
s = recv(1)
return '' if s.empty?
matching << s
sidx = matching.size
match_len = 0
delims.each do |d|
idx = matching.index(d)
next unless idx
if idx + d.size <= sidx + match_len
sidx = idx
match_len = d.size
end
end
if sidx < matching.size
r = data.get + matching.slice!(0, sidx + match_len)
r.slice!(-match_len..-1) if drop
return r
end
data << matching.slice!(0...-max_len) if matching.size > max_len
end
''
ensure
unrecv(matching)
unrecv(data)
end
end
end
# Receives a single line from the tube.
# A "line" is any sequence of bytes terminated by the byte sequence set in +context.newline+,
# which defaults to +"\\n"+.
#
# @!macro drop_definition
# @!macro timeout_definition
#
# @return [String]
# All bytes received over the tube until the first newline is received.
# Optionally retains the ending.
def recvline(drop: false, timeout: nil)
recvuntil(context.newline, drop: drop, timeout: timeout)
end
# Receives the next "line" from the tube; lines are separated by +sep+.
# The difference with +IO#gets+ is using +context.newline+ as default newline.
#
# @param [String, Integer] sep
# If +String+ is given, use +sep+ as the separator.
# If +Integer+ is given, receive exactly +sep+ bytes.
# @!macro drop_definition
# @!macro timeout_definition
#
# @return [String]
# The next "line".
#
# @raise [Pwnlib::Errors::EndOfTubeError]
# When the remaining data does not contain +sep+.
# When the size of the remaining data is less than +sep+.
#
# @example
# Sock.new('127.0.0.1', 1337).gets
# #=> "This is line one\n"
#
# Sock.new('127.0.0.1', 1337).gets(drop: true)
# #=> "This is line one"
#
# Sock.new('127.0.0.1', 1337).gets 'line'
# #=> "This is line"
#
# Sock.new('127.0.0.1', 1337).gets ''
# #=> "This is line"
#
# Sock.new('127.0.0.1', 1337).gets(4)
# #=> "This"
def gets(sep = context.newline, drop: false, timeout: nil)
case sep
when Integer
recvn(sep, timeout: timeout)
when String
recvuntil(sep, drop: drop, timeout: timeout)
else
raise ArgumentError, 'only Integer and String are supported'
end
end
# Wrapper around +recvpred+, which will return when a regex matches the string in the buffer.
#
# @param [Regexp] regex
# A regex to match.
# @!macro timeout_definition
#
# @return [String]
# A string contains bytes received from the tube, or +''+ if a timeout occurred while
# waiting.
def recvregex(regex, timeout: nil)
recvpred(timeout: timeout) { |data| data =~ regex }
end
# Receives data until reaching EOF or a timeout is occurred.
#
# @!macro timeout_definition
#
# @return [String]
# Returns the data received.
def recvall(timeout: nil)
recvn(1 << 63, timeout: timeout)
rescue ::Pwnlib::Errors::EndOfTubeError, ::Pwnlib::Errors::TimeoutError
@buffer.get
end
alias readall recvall
# Sends data.
#
# @param [String] data
# The +data+ string to be sent.
#
# @!macro send_return_definition
def send(data)
data = data.to_s
log.debug(format('Sent %#x bytes:', data.size))
log.indented(::Pwnlib::Util::HexDump.hexdump(data), level: DEBUG)
send_raw(data)
data.size
end
alias write send
# Sends the given object with +context.newline+.
#
# @param [Object] obj
# The object to be sent.
#
# @!macro send_return_definition
def sendline(obj)
s = obj.to_s + context.newline
write(s)
end
# Sends the given object(s).
# The difference with +IO#puts+ is using +context.newline+ as default newline.
#
# @param [Array<Object>] objs
# The objects to be sent.
#
# @!macro send_return_definition
#
# @example
# s.puts
# puts client.recv
# #
# #=> nil
#
# @example
# s.puts('shik', "hao\n", 123)
# puts client.recv
# # shik
# # hao
# # 123
# #=> nil
#
# @example
# s.puts(["darkhh\n\n", 'wei shi', 360])
# puts client.recv
# # darkhh
# #
# # wei shi
# # 360
# #=> nil
def puts(*objs)
return write(context.newline) if objs.empty?
objs = *objs.flatten
s = +''
objs.map(&:to_s).each do |elem|
s << elem
s << context.newline unless elem.end_with?(context.newline)
end
write(s)
end
# Does simultaneous reading and writing to the tube. In principle this just connects the tube
# to standard in and standard out.
def interact
log.info('Switching to interactive mode')
$stdout.write(@buffer.get)
until io_out.closed?
rs, = IO.select([$stdin, io_out])
if rs.include?($stdin)
s = $stdin.readpartial(BUFSIZE)
write(s)
end
if rs.include?(io_out)
s = recv
$stdout.write(s)
end
end
rescue ::Pwnlib::Errors::EndOfTubeError
log.info('Got EOF in interactive mode')
end
private
# Normalize direction.
#
# @param [Symbol] direction
#
# @return [Array<Symbol>]
# If +direction+ equals to
# * +:both+, returns +[:read, :write]+
# * +:read+ or +:recv+, returns [:read]
# * +:write+ or +:send+, returns [:write]
# Otherwise, raise +ArgumentError+.
def normalize_direction(direction)
case direction
when :both then %i[read write]
when :read, :recv then [:read]
when :write, :send then [:write]
else
raise ArgumentError, 'Only allow :both, :recv, :read, :send and :write passed'
end
end
def fillbuffer(timeout: nil)
data = @timer.countdown(timeout) do
self.timeout_raw = (@timer.timeout == :forever ? nil : @timer.timeout)
recv_raw(BUFSIZE)
end
if data
@buffer << data
log.debug(format('Received %#x bytes:', data.size))
log.indented(::Pwnlib::Util::HexDump.hexdump(data), level: DEBUG)
end
data
end
# The IO object of output, will be used for IO.select([io_out]) in interactive mode.
#
# @return [IO]
def io_out
raise NotImplementedError, 'Not implemented'
end
# @param [String] _data
#
# @return [void]
def send_raw(_data)
raise NotImplementedError, 'Not implemented'
end
# @param [Integer] _size
#
# @return [String]
def recv_raw(_size)
raise NotImplementedError, 'Not implemented'
end
# @param [Float?] _timeout
#
# @return [void]
def timeout_raw=(_timeout)
raise NotImplementedError, 'Not implemented'
end
include ::Pwnlib::Context
include ::Pwnlib::Logger
end
end
end