Skip to content
This repository
branch: master
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

executable file 208 lines (170 sloc) 3.82 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207
#! /usr/bin/env ruby
require 'eventmachine'
require 'thread'
require 'optparse'

begin
require 'readline'
require 'colorb'

module Readline
Commands = ['!exit', '!quit', '!clients']
Prefix = '>> '.bold

def self.supported?
true
end

self.completion_proc = proc {|s|
next unless s.start_with?('!')

Commands.grep(/^#{Regexp.escape(s)}/)
}
end
rescue LoadError
module Readline
def self.supported?
false
end
end
end

options = {}

OptionParser.new do |o|
options[:verbose] = false
options[:number] = 1000
options[:prefix] = 'S_'
options[:port] = 6667

o.on '-V', '--version' do
puts "Fail IRC stress test #{IRC.version}"
exit 0
end

o.on '-v', '--[no-]verbose' do |value|
options[:verbose] = value
end

o.on '-n', '--number VALUE', Integer do |value|
options[:number] = value.to_i
end

o.on '-s', '--server VALUE' do |value|
options[:server] = value
end

o.on '-p', '--port VALUE' do |value|
options[:port] = value.to_i
end

o.on '-P', '--prefix VALUE' do |value|
options[:prefix] = value
end
end.parse!

module Readline
def self.puts (text)
print "\r#{text}#{' ' * (Readline.get_screen_size.last - text.length - 1)}\n"
print Prefix
end

def self.readline_with_hist_management
begin
line = Readline.readline(Prefix, true)
rescue Exception
return
end

return unless line

if line =~ /^\s*$/ or Readline::HISTORY.to_a[-2] == line
Readline::HISTORY.pop
end

line
end
end

class StressBot < EventMachine::Connection
include EventMachine::Protocols::LineText2

def initialize(host, port, name, channel)
@host, @port, @name, @channel = host, port, name, channel
@closing = false

@sid = channel.subscribe {|msg|
send_data(msg)
}
end

def connection_completed
send_data "NICK #{@name}\r\nUSER #{@name} #{@name} #{@name} :#{@name}\r\n"
end

def receive_line(line)
if line =~ /^PI(NG.*)$/
send_data("PO#{$1}\r\n")
end
end

def game_over
@closing = true
send_data("QUIT :GAME OVER\r\n")
@channel.unsubscribe(@sid)
close_connection
end

def reconnect
EventMachine.reconnect @host, @port, self
end

def unbind
EventMachine.add_timer(5) do
reconnect
end unless @closing
end
end

class Stresser
def initialize(host, port, prefix, number)
@host, @port, @prefix, @number = host, port, prefix, number
@channel = EventMachine::Channel.new
end

def start
blk = proc {
@clients = (1..@number).map {|no|
EventMachine.connect @host, @port, StressBot, @host, @port, "#{@prefix}#{no}", @channel
}
}

EventMachine.reactor_running? ? blk.call : EventMachine.run(&blk)
end

def number=(num)
return if num == @clients.size

if num < @clients.size
(num ... @clients.size).to_a.reverse.each do |i|
@clients[i].game_over
@clients.delete(i)
end
else
(@clients.size ... num).each do |i|
EventMachine.schedule do
p i
@clients[i] = EventMachine.connect @host, @port, StressBot, @host, @port, "#{@prefix}#{i + 1}", @channel
end
end
end
end

def puts(data)
@channel.push("#{data.gsub(/\r?\n$/, '')}\r\n")
end

def close
@clients.each(&:game_over)
EventMachine.stop_event_loop
end
end

stresser = nil
Thread.abort_on_exception = true
thread = Thread.new do
(stresser = Stresser.new(options[:server], options[:port], options[:prefix], options[:number])).start
end

['TERM', 'KILL', 'INT'].each do |sig|
trap sig do
stresser.close
thread.join
exit 0
end
end

if Readline.supported?
while line = Readline.readline_with_hist_management
if line.start_with?('!')
case line[1..-1]
when 'exit', 'quit'
stresser.close
thread.join
exit!
when /^clients\s*(.*)$/
number = $1.to_i
next if number.zero?
stresser.number = number
end
else
stresser.puts(line)
end
end
else
thread.join
end

Something went wrong with that request. Please try again.