Socket Streamer

req_file

danger will robinson

streamer.rb

#!/usr/bin/env ruby
require 'socket'
require 'configliere'

Settings.use :commandline

Settings.define :rate, type: Float, default: (1.0/0.0), description: "Max request rate in reqs/sec. Defaults to unbound."
Settings.define :req_file, default: 'req_file', description: "File to grab requests from. Streams random characters by default."
Settings.define :silent, default: false, type: :boolean, description: "Stay silent"
Settings.define :reopen_socket, default: false, type: :boolean, description: "Reopen socket when the end of req_file is reached"
Settings.define :udp, default:false, type: :boolean, description: "Open a UDP socket instead of a TCP socket"
Settings.define :syslog, default: false, type: :boolean, description: "Wrap requests in a syslog format"

Settings.resolve!

host = 'localhost'
port = '9000'

if Settings.rest.length > 0
  url = Settings.rest[0].split(':')
  host = url[0..-2].join
  port = url[-1]
end

if Settings.udp
  t = UDPSocket.new
  t.connect(host,port)
else
  t = TCPSocket.new(host, port)
end

req_src = ($stdin.tty?) ? File.open(Settings.req_file,'r') : $stdin

reqs_remaining = Settings.rate
next_t = Time.now.to_f + 1.0
puts "Streaming to #{host}:#{port} at #{Settings.rate} reqs/sec"

loop do
  req_src.each_line do |req|
    now = Time.now.to_f
    if Settings.syslog
      time = Time.now.strftime("%b %d %H:%M:%S 10.78.7.215 foo:")
      req = "<12> #{time} #{req}"
    end
    if now >= next_t or reqs_remaining <= 0
      next_t += 1.0
      reqs_remaining += Settings.rate
    end
    time_remaining = next_t - now
    t.puts req
    reqs_remaining -= 1
    sleep(time_remaining/reqs_remaining)
  end
  req_src.rewind if STDIN.tty?
  if Settings.reopen_socket
    t.close
    if Settings.udp
      t = UDPSocket.new
      t.connect(host,port)
    else
      t = TCPSocket.new(host,port) 
    end
  end
end

tcp_server.rb

#!/usr/bin/env ruby
require 'socket'
require 'configliere'

Settings.use :commandline

Settings.define :port, default: 9000
Settings.define :udp, default: false, type: :boolean
Settings.resolve!

t = TCPServer.new Settings.port
puts "Listening on port #{Settings.port}"

loop do  
  client = t.accept
  puts "Client connected on port #{Settings.port}"
  then_t = Time.now.to_f
  num_reqs = 0
  while line = client.gets
    num_reqs += 1
    now_t = Time.now.to_i
    if now_t - then_t >= 1.0
      puts "%.2f reqs/sec \n" % (num_reqs/(now_t - then_t))
      num_reqs = 0
      then_t = Time.now.to_f
    end
  end
end
s.close

udp_server.rb

#!/usr/bin/env ruby
require 'socket'
require 'configliere'

Settings.use :commandline

Settings.define :port, default: 9000
Settings.define :udp, default: false, type: :boolean
Settings.define :maxlen, default: 65536, type: Integer
Settings.resolve!

t = UDPSocket.new
t.bind(nil,Settings.port)

puts "Listening on port #{Settings.port}"

num_reqs = 0
then_t = Time.now.to_f
loop do
  t.recvfrom(Settings.maxlen)
  num_reqs += 1
  now_t = Time.now.to_i
  if now_t - then_t >= 1.0
    puts "%.2f reqs/sec \n" % (num_reqs/(now_t - then_t))
    num_reqs = 0
    then_t = Time.now.to_f
  end
end
s.close