diff --git a/app/controllers/vuln_api.rb b/app/controllers/vuln_api.rb index a656504c..605f0ccc 100644 --- a/app/controllers/vuln_api.rb +++ b/app/controllers/vuln_api.rb @@ -4,6 +4,8 @@ module WPScan module Controller # Controller to handle the API token class VulnApi < CMSScanner::Controller::Base + ENV_KEY = 'WPSCAN_API_TOKEN' + def cli_options [ OptString.new(['--api-token TOKEN', 'The WPVulnDB API Token to display vulnerability data']) @@ -11,9 +13,9 @@ module WPScan end def before_scan - return unless ParsedCli.api_token + return unless ParsedCli.api_token || ENV.key?(ENV_KEY) - DB::VulnApi.token = ParsedCli.api_token + DB::VulnApi.token = ParsedCli.api_token || ENV[ENV_KEY] api_status = DB::VulnApi.status diff --git a/spec/app/controllers/vuln_api_spec.rb b/spec/app/controllers/vuln_api_spec.rb index f2714c49..dcae474b 100644 --- a/spec/app/controllers/vuln_api_spec.rb +++ b/spec/app/controllers/vuln_api_spec.rb @@ -74,20 +74,40 @@ describe WPScan::Controller::VulnApi do context 'when limited requests' do let(:requests) { 100 } - it 'does not raise an error' do + it 'sets the token and does not raise an error' do expect { controller.before_scan }.to_not raise_error + + expect(WPScan::DB::VulnApi.token).to eql 'token' end context 'when unlimited requests' do let(:requests) { 'Unlimited' } - it 'does not raise an error' do + it 'sets the token and does not raise an error' do expect { controller.before_scan }.to_not raise_error + + expect(WPScan::DB::VulnApi.token).to eql 'token' end end end end end end + + context 'when token in ENV' do + before do + ENV[described_class::ENV_KEY] = 'token-from-env' + + expect(WPScan::DB::VulnApi) + .to receive(:status) + .and_return('success' => true, 'plan' => 'free', 'requests_remaining' => 'Unlimited') + end + + it 'sets the token and does not raise an error' do + expect { controller.before_scan }.to_not raise_error + + expect(WPScan::DB::VulnApi.token).to eql 'token-from-env' + end + end end end