wasp 0.2.0

data/Gemfile

@@ -0,0 +1,5 @@
+source "http://rubygems.org"
+
+gem "aws-sdk", ">= 1.3.2"
+gem "net-ssh", ">= 2.3.0"
+gem "net-scp", ">= 1.0.4"

data/LICENSE

@@ -0,0 +1,20 @@
+Copyright (c) 2012 Ji-Hun Seol
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/README.md

@@ -0,0 +1,77 @@
+WaspswithRainofStings
+====================
+
+By arkemiq (arkemiq@gmail.com)
+
+Description
+-----------
+
+Wasp is a distributed heavy load generator in Ruby inspired by beeswithmachineguns(https://github.com/newsapps/beeswithmachineguns).
+
+This is a utility of arming(create and install load program) many wasps (micro EC2 instances)
+from multiple regions(EC2 zones) to attack (load test) targets(web applications).
+
+Installation
+------------
+
+	$ git clone https://github.com/arkemiq/wasp.git
+	$ cd wasp
+	$ bundle install
+
+Configuring EC2 credentials
+---------------------------
+	
+Wasps uses aws-sdk gem to communicate with EC2 and supports all the same methods of storing credentials that it does.
+Set access key and secret key for config/aws.yml to access your EC2 account.
+
+	access_key_id: xxxxxxxxxxxxxxxxx
+	secret_access_key: xxxxxxxxxxxxxxxxx
+
+Running
+-------
+
+	Usage:
+	
+	To set credentials:
+	  $ wasp set [AWS Credential file]
+	
+	To launch 6 wasps:
+	  - launch 6 instances in us-east zone with private key named wasps
+	  $ wasp up -k wasps -s 6 
+	
+	  - launch 5 instances in us-west-2 zone with ami-8cb33ebc AMI, username ubuntu and private key named wasps
+	  $ wasp up -k wasps -z us-west-2 -a ami-8cb33ebc -s 5 -l ubuntu
+	
+	To check status:
+	  $	wasp status
+	
+	To equip weapon(default: apachebench):
+	  $ wasp equip
+
+	To attack target with 1000 requests and 100 concurrent users per each wasp:
+	  $ wasp attack -n 1000 -c 100 -u http://target_site
+
+    To attack target with incrementally increase wasps from 1 to 10000 during 60 seconds:
+	  $ wasp rattack -p 10000:60 -u http://target_site
+	
+	To sleep wasps:
+	  $ wasp down
+	
+	To see options:
+	  $ wasp help
+	
+*Note*: the default EC2 security group is called 'wasps' and by default it locks out SSH access.
+
+*Note*: if you didn't specify a zone for EC2 instances, it use default zone 'us-east'.
+
+*Please remember to do this*--we aren't responsible for your EC2 bills.
+
+Notice! (PLEASE READ)
+---------------------
+If you decide to use the Wasps, please keep in mind the following important notice: they are, more-or-less a distributed denial-of-service attack in a fancy package and, therefore, if you point them at any server you don’t own you will behaving *unethically*, have your Amazon Web Services account *locked-out*, and be *liable* in a court of law for any downtime you cause.
+
+You have waspn warned.
+
+License
+-------
+MIT.

data/Rakefile

@@ -0,0 +1,15 @@
+require 'rake/testtask'
+
+desc "Run rspec"
+task :spec do
+  sh('bundle install')
+  require "rspec/core/rake_task"
+  RSpec::Core::RakeTask.new do |t|
+    t.rspec_opts = %w(-fs -c)
+  end
+end
+task :default => :spec
+
+Rake::TestTask.new do |t|
+  t.libs << 'test'
+end

data/bin/wasp

@@ -0,0 +1,6 @@
+#!/usr/bin/env ruby
+
+require File.expand_path('../../lib/wasp', __FILE__)
+
+WASP::Nest.wakeup(ARGV.dup)
+

data/config/aws.yml

@@ -0,0 +1,4 @@
+# Fill in your AWS Access Key ID and Secret Access Key
+# http://aws.amazon.com/security-credentials
+access_key_id: 
+secret_access_key: 

data/lib/wasp.rb

@@ -0,0 +1,14 @@
+ROOT = File.expand_path(File.dirname(__FILE__))
+
+module WASP
+  autoload :Const,    	"#{ROOT}/wasp/const"
+  autoload :Wasp,      	"#{ROOT}/wasp/wasp"
+  autoload :QueenWasp, 	"#{ROOT}/wasp/queenwasp"
+  autoload :Nest,  		"#{ROOT}/wasp/nest"
+  autoload :Aws,      	"#{ROOT}/wasp/ec2"
+  autoload :Config,   	"#{ROOT}/wasp/config"
+  autoload :WeaponAB, 	"#{ROOT}/wasp/stingab"
+end
+
+require "#{ROOT}/wasp/core_ext"
+require 'fileutils'

data/lib/wasp/config.rb

@@ -0,0 +1,11 @@
+
+module WASP
+  class Config
+    
+    class << self
+      attr_accessor :output
+      attr_accessor :colorize
+    end
+    
+  end
+end

data/lib/wasp/const.rb

@@ -0,0 +1,21 @@
+module WASP
+  module Const
+    VERSION = '0.2.0'
+  
+    DEFAULT_WASPS = 5
+    DEFAULT_GROUP = 'wasps'
+    DEFAULT_ZONE = 'us-east-1'
+    DEFAULT_AMI = 'ami-baba68d3'
+    DEFAULT_USER = 'ubuntu'
+    DEFAULT_KEY = 'wasp'
+    DEFAULT_NUMBER_OF_REQUESTS = 1000
+    DEFAULT_TIME = 10
+    DEFAULT_CONCURRENT_OF_CONNECTIONS = 100
+    DEFAULT_WEAPON = 'ab'
+    SSH_PORT = 22
+    NATS_PORT = 4222
+    DEFAULT_MAXIMUM_STINGS = 10000
+    DEFAULT_WAVE_TIME = 10
+    DEFAULT_COOLDOWN = 10
+  end
+end

data/lib/wasp/core_ext.rb

@@ -0,0 +1,113 @@
+module WASPExtensions
+
+  def say(message)
+    WASP::Config.output.puts(message) if WASP::Config.output
+  end
+
+  def header(message, filler = '-')
+    say "\n"
+    say message
+    say filler.to_s * message.size
+  end
+
+  def banner(message)
+    say "\n"
+    say message
+  end
+
+  def display(message, nl=true)
+    if nl
+      say message
+    else
+      if WASP::Config.output
+        WASP::Config.output.print(message)
+        WASP::Config.output.flush
+      end
+    end
+  end
+
+  def clear(size=80)
+    return unless WASP::Config.output
+    WASP::Config.output.print("\r")
+    WASP::Config.output.print(" " * size)
+    WASP::Config.output.print("\r")
+    #WASP::Config.output.flush
+  end
+
+  def err(message, prefix='Error: ')
+    raise WASP::CliExit, "#{prefix}#{message}"
+  end
+
+  def quit(message = nil)
+    raise WASP::GracefulExit, message
+  end
+
+  def blank?
+    self.to_s.blank?
+  end
+
+  def uptime_string(delta)
+    num_seconds = delta.to_i
+    days = num_seconds / (60 * 60 * 24);
+    num_seconds -= days * (60 * 60 * 24);
+    hours = num_seconds / (60 * 60);
+    num_seconds -= hours * (60 * 60);
+    minutes = num_seconds / 60;
+    num_seconds -= minutes * 60;
+    "#{days}d:#{hours}h:#{minutes}m:#{num_seconds}s"
+  end
+
+  def pretty_size(size, prec=1)
+    return 'NA' unless size
+    return "#{size}B" if size < 1024
+    return sprintf("%.#{prec}fK", size/1024.0) if size < (1024*1024)
+    return sprintf("%.#{prec}fM", size/(1024.0*1024.0)) if size < (1024*1024*1024)
+    return sprintf("%.#{prec}fG", size/(1024.0*1024.0*1024.0))
+  end
+
+end
+
+module StringExtensions
+    def red
+        colorize("\e[0m\e[31m")
+    end
+
+    def green
+        colorize("\e[0m\e[32m")
+    end
+
+    def yellow
+        colorize("\e[0m\e[33m")
+    end
+
+    def bold
+        colorize("\e[0m\e[1m")
+    end
+
+    def colorize(color_code)
+        "#{color_code}#{self}\e[0m"
+    end
+
+    def blank?
+        self =~ /^\s*$/
+    end
+
+    def truncate(limit = 30)
+        return "" if self.blank?
+        etc = "..."
+        stripped = self.strip[0..limit]
+        if stripped.length > limit
+            stripped.gsub(/\s+?(\S+)?$/, "") + etc
+        else
+            stripped
+        end
+    end
+end
+
+class Object
+  include WASPExtensions
+end
+
+class String
+  include StringExtensions
+end

data/lib/wasp/ec2.rb

@@ -0,0 +1,355 @@
+require 'aws-sdk'
+
+module WASP
+  class Aws
+    attr_reader   :config_path
+    attr_reader   :ec2
+    attr_reader   :group
+    attr_reader   :key
+    attr_reader   :instances
+    attr_reader   :regions
+    attr_reader   :instance_list
+    attr_reader   :threads
+    attr_reader   :login
+  
+    SLEEP_TIME = 1
+    LINE_LENGTH = 100
+  
+    def initialize (args)
+      config_path = ENV["HOME"] + "/.waspaws.yml"
+      begin
+        AWS.config(YAML.load(File.read(config_path)))
+        @ec2 = AWS::EC2.new
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+        puts "[WARN]".yellow + " Please set AWS credential file."
+        exit false
+      end
+      
+      # evaluate AWS access_key and secret_access_key
+      begin
+        print "EC2".green + " Checking access key validation.."
+        @ec2.availability_zones.each do |av| 
+          av.name 
+        end
+        puts " OK".green
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+        puts "[WARN]".yellow + " Please copy/paste correct AWS access_key and secret_access_key to config/aws.yml file"
+        exit false
+      end
+      
+      @num_wasps = if args[:server].nil? then WASP::Const::DEFAULT_WASPS
+                  else
+                    args[:server].to_i
+                  end
+      @group = if args[:group].nil? then WASP::Const::DEFAULT_GROUP
+               else
+                 args[:group]
+               end
+      @zone = if args[:zone].nil? then WASP::Const::DEFAULT_ZONE
+              else
+                args[:zone]
+              end              
+      @ami = if args[:ami].nil? then get_default_ami(@zone)
+             else
+               args[:ami]
+             end
+      @login = if args[:login].nil? then WASP::Const::DEFAULT_USER
+               else
+                 args[:login]
+               end     
+      
+      @key = args[:key]
+      
+      
+      @regions = @ec2.regions
+      @instance_list = []
+      @instances = nil
+    end
+    
+    def get_keypair 
+      make_key = true
+      key = @regions[@zone].key_pairs.filter('key-name', @key + '-' + @zone).first
+      
+      if key.nil? == false then
+        if not File.exists?("#{ENV['HOME']}/.ssh/#{@key}-#{@zone}.pem") then
+          key.delete 
+          make_key = true
+        else
+          make_key = false
+        end
+      end
+                  
+      if make_key then
+        key = @regions[@zone].key_pairs.create(@key + '-' + @zone)
+        keyfile = "#{ENV['HOME']}/.ssh/#{@key}-#{@zone}.pem"
+        begin
+          File.delete(keyfile) if File.exists?(keyfile)
+          File.open(keyfile, "w") do |f|
+            f.write(key.private_key)
+          end
+          File.chmod(0400, keyfile)
+        rescue => ex
+          print "[WARN]".yellow + " #{ex.message}"
+          exit false
+        end
+      end
+
+      puts "EC2".green + " Private key is created in " + "~/.ssh/#{@key}-#{@zone}.pem".bold
+    end
+    
+    def get_security_group
+      found = false
+      puts "EC2".green + " Look up existed security group [" + "#{@group}".bold + "].."
+      security_groups = @regions[@zone].security_groups
+      
+      found = true if security_groups.filter('group-name', @group).first != nil
+      
+      if not found then
+        puts "EC2".green + " Creating security group " + "#{@group}".bold + ".."
+        wasps = security_groups.create(@group)
+        puts "EC2".green + " Open inbound tcp port :#{WASP::Const::SSH_PORT}"
+        wasps.authorize_ingress(:tcp, WASP::Const::SSH_PORT)
+        puts "EC2".green + " Open inbound tcp port :#{WASP::Const::NATS_PORT}"
+        wasps.authorize_ingress(:tcp, WASP::Const::NATS_PORT)
+        puts "EC2".green + " Allow ping"
+        wasps.allow_ping
+      end
+      
+      puts "EC2".green + " Security group [#{@group}] have set.."
+    end
+    
+    def waiting_wasp (i)
+      banner = "EC2".green + " Launching wasps: "
+      
+      display banner, false      
+      while i.status == :pending do
+        print '.'
+        sleep SLEEP_TIME
+      end
+    
+      if i.status == :running then
+        clear(LINE_LENGTH)
+        display "#{banner}#{'OK'.green}"
+        @instance_list.push(i.id)
+        puts "Wasp " + "#{i.id}".yellow + " is ready to attack"
+        i.tags.Name = 'a wasp!'
+      else
+        puts "EC2".green + " #{i.id} is going to wrong place"
+      end
+    end
+  
+    def create 
+      puts "EC2".green + " #{@num_wasps} wasps will be launched.."
+      
+      begin
+        @instances = @ec2.regions[@zone].instances.create(:image_id=>@ami,
+                                              :security_groups=>[@group],
+                                              :key_name=>@key + '-' + @zone,
+                                              :instance_type=>'t1.micro',
+                                              :count=>@num_wasps)                                      
+        num_wasps = 0
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+        return false
+      end
+      if @instances.class == Array then
+        @instances.each do |i|
+          waiting_wasp(i)
+        end
+        num_wasps = @instances.count
+      else
+        waiting_wasp(@instances)
+        num_wasps = 1
+      end
+      
+      _write_to_file
+    
+      puts "The swarm has assembled " + "#{num_wasps}".yellow + " wasps.\n"
+      true
+    end
+    
+    def _get_instances
+      status_of_wasps = {}
+      if @instance_list.count == 0 then
+        begin
+          File.open("#{ENV["HOME"]}/.nest", "r") do |file|
+            lines = file.readlines
+            lines.each do |line|
+              instance, key, zone, login = line.split(' ')
+              
+              if status_of_wasps[zone].nil? then
+                status_of_wasps[zone] = {}
+                status_of_wasps[zone][:instances] = []
+                status_of_wasps[zone][:instances].push(instance)
+                status_of_wasps[zone][:login] = login
+              else
+                status_of_wasps[zone][:instances].push(instance)
+              end
+            end
+          end
+        rescue => ex
+          puts "There is no wasps in the air."
+          
+          return nil
+        end
+      end
+
+      return status_of_wasps
+    end
+    
+    def _write_to_file
+      begin
+        File.open("#{ENV["HOME"]}/.nest", "a+") do |file|
+          if @instances.class == Array then
+            @instances.each do |i|
+              file.puts("#{i.id} #{i.key_name} #{i.availability_zone} #{@login}\n")
+            end
+          else
+            file.puts("#{@instances.id} #{@instances.key_name} #{@instances.availability_zone} #{@login}\n")
+          end
+        end
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+      end
+    end
+    
+    def _delete_file
+      begin
+        File.delete("#{ENV["HOME"]}/.nest")
+      rescue => ex
+        puts "[WARN]".yellow + "#{ex.message}"
+        puts "There is no wasps in the air."
+      end
+    end
+  
+    def terminate
+      count = 0
+      
+      #if @instances.nil? then
+        ins = _get_instances
+        
+        return nil if ins.nil?
+        
+        
+        ins.each do |region, info|
+          r = get_ec2_zone(region)
+          info[:instances].each do |i| 
+            count += 1
+            puts "Wasp" + " #{i}".yellow + " from #{region} is going home"
+            begin
+              @ec2.regions[r].instances[i].terminate
+            rescue => e
+              puts "#{e.message}"
+            end
+          end
+        end
+      
+      _delete_file
+      
+      puts "EC2".green + " Stood down #{count} wasps.\n"
+    end
+    
+    def get_wasps
+      wasps = []
+      ins = _get_instances
+      
+      return nil if ins.nil?
+      
+      count = ins.count
+      ins.each do |region, info|
+        r = get_ec2_zone(region)
+        login = info[:login]
+        info[:instances].each do |it|
+          i = @ec2.regions[r].instances[it]
+        
+          wasps.push({:instance_id => i.id, 
+                     :instance_name => i.dns_name, 
+                     :region => region,
+                     :login => login,
+                     :key_name => i.key_pair.name,
+                     :report => nil,
+                     :wavereport => nil})
+        end
+      end
+      
+      wasps
+    end
+    
+    def status
+      return ins = _get_instances
+    end
+    
+    def get_wasp_status(zone=WASP::Const::DEFAULT_ZONE , id)
+      begin
+        z = get_ec2_zone(zone)
+        @ec2.regions[z].instances[id].status
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+        return nil
+      end        
+    end
+    
+    def get_wasp_domain(zone=WASP::Const::DEFAULT_ZONE , id)
+      begin
+        z = get_ec2_zone(zone)
+        @ec2.regions[z].instances[id].dns_name
+      rescue => ex
+        puts "[WARN]".yellow + " #{ex.message}"
+        return nil
+      end
+    end
+    
+    def get_ec2_zone(zone)
+      case
+      when zone.match(/eu-west-1/)
+        "eu-west-1"
+      when zone.match(/sa-east-1/)
+        "sa-east-1"
+      when zone.match(/us-west-1/)
+        "us-west-1"
+      when zone.match(/us-west-2/)
+        "us-west-2"
+      when zone.match(/ap-northeast-1/)
+        "ap-northeast-1"
+      when zone.match(/us-east-1/)
+        "us-east-1"
+      when zone.match(/ap-southeast-1/)
+        "ap-southeast-1"
+      else
+        "us-east-1"
+      end
+    end
+    
+    def get_default_ami(zone)
+      case
+      when zone.match(/eu-west-1/)
+        "ami-895069fd"
+      when zone.match(/sa-east-1/)
+        "ami-b673acab"
+      when zone.match(/us-west-1/)
+        "ami-6da8f128"
+      when zone.match(/us-west-2/)
+        "ami-ae05889e"
+      when zone.match(/ap-northeast-1/)
+        "ami-10299f11"
+      when zone.match(/us-east-1/)
+        "ami-baba68d3"
+      when zone.match(/ap-southeast-1/)
+        "ami-4296d210"
+      else
+        "ami-baba68d3"
+      end
+    end
+  
+    def get_regions
+      regions = []
+      @ec2.regions.each do |r|
+        regions.push(r.name)
+      end
+      
+      regions
+    end
+  end
+end