Add a support for a custom user agent.

bin/concuss

@@ -19,6 +19,10 @@ OptionParser.new do |opts|
   opts.on("-t", "--test-string STRING", "Set a custom test string. If none specified, it sets a random string to match on.") do |v|
     options[:test_string] = v
   end
+
+  opts.on("-a", "--user-agent STRING", "Set a custom user agent. If none specified, it defaults to Concuss/#{Concuss::VERSION}") do |v|
+    options[:user_agent] = v
+  end
 end.parse!
 
 if ARGV[0].nil?

lib/concuss.rb

@@ -3,19 +3,20 @@
 class Concuss
   class Error < StandardError; end
 
-  attr_reader :url, :file, :header_set, :headers, :test_string
+  attr_reader :url, :file, :header_set, :headers, :test_string, :user_agent
 
-  def initialize(url:, file: nil, header_set: :all, test_string: nil)
+  def initialize(url:, file: nil, header_set: :all, test_string: nil, user_agent: nil)
     @url = url
     @file = file
     @header_set = file.nil? ? header_set : :file
     @test_string = test_string
+    @user_agent = user_agent
 
     @headers = Concuss::Headers.new(header_set: @header_set, file: @file).group
   end
 
   def attack!
-    runner = Concuss::Runner.new(headers: headers, url: url, test_string: test_string)
+    runner = Concuss::Runner.new(headers: headers, url: url, test_string: test_string, user_agent: user_agent)
 
     runner.run
   end

lib/concuss/headers.rb

@@ -30,7 +30,6 @@ class Concuss::Headers
     'Range',
     'Referer',
     'TE',
-    'User-Agent',
     'Upgrade',
     'Via',
     'Warning'

lib/concuss/runner.rb

@@ -2,20 +2,26 @@
 require 'net/http'
 
 class Concuss::Runner
-  attr_reader :headers, :url, :test_string
+  DEFAULT_USER_AGENT = "Concuss/#{Concuss::VERSION}"
+
+  attr_reader :headers, :url, :test_string, :user_agent
 
-  def initialize(headers:, url:, test_string: nil)
+  def initialize(headers:, url:, test_string: nil, user_agent: nil)
     @headers = headers
     @url = url
     @test_string = test_string || SecureRandom.hex(25)
+    @user_agent = user_agent || DEFAULT_USER_AGENT
   end
 
   def run
     uri = URI(@url)
 
     @headers.each do |header|
       response = Net::HTTP.get_response(uri,
-        { header => test_string }
+        {
+          header => test_string,
+          'User-Agent' => user_agent
+        }
       )
 
       if response.code == "200" && response.body.include?(@test_string)

spec/concuss/runner_spec.rb

@@ -4,7 +4,8 @@
   let(:headers) { ['X-Test-Header1', 'X-Test-Header2'] }
   let(:url) { 'http://test.com' }
   let(:test_string) { 'teststring' }
-  let(:runner) { Concuss::Runner.new(headers: headers, url: url, test_string: test_string) }
+  let(:runner) { Concuss::Runner.new(headers: headers, url: url, test_string: test_string, user_agent: user_agent) }
+  let(:user_agent) { 'Mozilla/5.0' }
 
   describe '#initialize' do
     it 'sets the headers' do
@@ -18,23 +19,50 @@
     it 'sets the test string' do
       expect(runner.test_string).to eq(test_string)
     end
+
+    it 'sets the user_agent' do
+      expect(runner.user_agent).to eq(user_agent)
+    end
   end
 
   describe '#run' do
+    let(:mock_response) { double(code: '200', body: test_string) }
+
     before(:each) do
       # mock the request to return a 200 status code
-      allow(Net::HTTP).to receive(:get_response).and_return(double(code: '200', body: test_string))
+      allow(Net::HTTP).to receive(:get_response).and_return(mock_response)
     end
 
     it 'sends a request for each of the headers' do
-      allow(Net::HTTP).to receive(:get_response).and_return(double(code: '200', body: test_string)).twice
+      allow(Net::HTTP).to receive(:get_response).and_return(mock_response).twice
       runner.run
     end
 
     it 'outputs the header, status code, and result' do
       expect { runner.run }.to output("X-Test-Header1 - 200 - HIT\nX-Test-Header2 - 200 - HIT\n").to_stdout
     end
 
+    it 'uses the custom user_agent' do
+      expect(Net::HTTP).to receive(:get_response).with(URI(url),
+        {
+          'User-Agent' => user_agent,
+          headers.first => test_string
+        }
+      )
+      runner.run
+    end
+
+    it 'uses the DEFAULT_USER_AGENT' do
+      runner = Concuss::Runner.new(headers: headers, url: url, test_string: test_string)
+      expect(Net::HTTP).to receive(:get_response).with(URI(url),
+        {
+          'User-Agent' => Concuss::Runner::DEFAULT_USER_AGENT,
+          headers.first => test_string
+        }
+      )
+      runner.run
+    end
+
     context 'when the response body does not include the test string' do
       before(:each) do
         # mock the request to return a 200 status code and body without the test string

spec/concuss_spec.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 RSpec.describe Concuss do
+  let(:user_agent) { 'Mozilla/5.0' }
+
   it "has a version number" do
     expect(Concuss::VERSION).not_to be nil
   end
@@ -18,14 +20,16 @@
       concuss = Concuss.new(
         url: 'https://example.com',
         file: file_path,
-        test_string: 'example'
+        test_string: 'example',
+        user_agent: user_agent
       )
 
       expect(concuss.url).to eq('https://example.com')
       expect(concuss.file).to eq(file_path)
       expect(concuss.header_set).to eq(:file)
       expect(concuss.test_string).to eq('example')
       expect(concuss.headers).to eq(fake_headers)
+      expect(concuss.user_agent).to eq(user_agent)
     end
 
     it 'sets a default header_set if none are passed in' do
@@ -41,10 +45,15 @@
 
   describe '#attack!' do
     it 'creates a new instance of Concuss::Runner' do
-      concuss = Concuss.new(url: 'https://example.com')
+      concuss = Concuss.new(url: 'https://example.com', user_agent: user_agent)
       runner = double('runner')
       expect(runner).to receive(:run)
-      expect(Concuss::Runner).to receive(:new).with(url: 'https://example.com', headers: Concuss::Headers.new(header_set: :all).group, test_string: nil).and_return(runner)
+      expect(Concuss::Runner).to receive(:new).with(
+        url: 'https://example.com',
+        user_agent: user_agent,
+        headers: Concuss::Headers.new(header_set: :all).group,
+        test_string: nil
+      ).and_return(runner)
       concuss.attack!
     end
   end