lib/shuttle/runner.rb in shuttle-deploy-0.2.0.beta2 vs lib/shuttle/runner.rb in shuttle-deploy-0.2.0.beta3

@@ -79,21 +79,25 @@
       end
 
       ssh.open
 
       klass = Shuttle.const_get(strategy.capitalize) rescue nil
+      command.gsub!(/:/,'_')
+      exit_code = 0
 
       if klass.nil?
         STDERR.puts "Invalid strategy: #{strategy}"
         exit 1
       end
 
-      integration = klass.new(config, ssh, server, target) 
+      unless %w(setup deploy rollback).include?(command)
+        STDERR.puts "Invalid command: #{command}"
+        exit 1
+      end
 
-      command.gsub!(/:/,'_')
-      exit_code = 0
-      puts "\n"
+      integration = klass.new(config, ssh, server, target)
 
+      puts "\n"
       puts "Shuttle v#{Shuttle::VERSION}\n"
       puts "\n"
       integration.log "Connected to #{server.user}@#{server.host}"
       
       if integration.respond_to?(command)
\ No newline at end of file