From ebc65160957888e5407d955b82a2c6c904fb6e0a Mon Sep 17 00:00:00 2001
From: Joshua Frankel <joshmfrankel@gmail.com>
Date: Fri, 8 May 2020 14:49:45 -0400
Subject: [PATCH] Add criticality filtering to cli

Why?

When utilizing bundler-audit on CI, it can be helpful
to filter criticality to find higher criticality gems.

What?

* Adds filter (--filter or -f) command line argument that
accepts array of values
* Filters within Audit::Scanner to remove filtered out advisories
---
 README.md                    |  4 ++++
 lib/bundler/audit/cli.rb     |  3 ++-
 lib/bundler/audit/scanner.rb | 16 ++++++++++++----
 spec/integration_spec.rb     | 18 ++++++++++++++++++
 spec/scanner_spec.rb         | 10 ++++++++++
 5 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 7c8b5e6d..d53f5fa8 100644
--- a/README.md
+++ b/README.md
@@ -116,6 +116,10 @@ Ignore specific advisories:
 
     $ bundle-audit check --ignore OSVDB-108664
 
+Filter results by criticality level (low, medium, high):
+
+    $ bundle-audit check --filter medium high
+
 Rake task:
 
 ```ruby
diff --git a/lib/bundler/audit/cli.rb b/lib/bundler/audit/cli.rb
index adbb7e6b..d9f7a325 100644
--- a/lib/bundler/audit/cli.rb
+++ b/lib/bundler/audit/cli.rb
@@ -34,6 +34,7 @@ class CLI < ::Thor
       method_option :verbose, :type => :boolean, :aliases => '-v'
       method_option :ignore, :type => :array, :aliases => '-i'
       method_option :update, :type => :boolean, :aliases => '-u'
+      method_option :filter, :type => :array, :aliases => '-f'
 
       def check
         update if options[:update]
@@ -41,7 +42,7 @@ def check
         scanner    = Scanner.new
         vulnerable = false
 
-        scanner.scan(:ignore => options.ignore) do |result|
+        scanner.scan(ignore: options[:ignore], filter: options[:filter]) do |result|
           vulnerable = true
 
           case result
diff --git a/lib/bundler/audit/scanner.rb b/lib/bundler/audit/scanner.rb
index eb277dc8..1597f78f 100644
--- a/lib/bundler/audit/scanner.rb
+++ b/lib/bundler/audit/scanner.rb
@@ -56,6 +56,9 @@ def initialize(root=Dir.pwd,gemfile_lock='Gemfile.lock')
       # @option options [Array<String>] :ignore
       #   The advisories to ignore.
       #
+      # @option options [Array<String>] :filter
+      #   The criticalities to filter.
+      #
       # @yield [result]
       #   The given block will be passed the results of the scan.
       #
@@ -68,13 +71,10 @@ def initialize(root=Dir.pwd,gemfile_lock='Gemfile.lock')
       def scan(options={},&block)
         return enum_for(__method__,options) unless block
 
-        ignore = Set[]
-        ignore += options[:ignore] if options[:ignore]
-
         scan_sources(options,&block)
         scan_specs(options,&block)
 
-        return self
+        self
       end
 
       #
@@ -127,6 +127,9 @@ def scan_sources(options={})
       # @option options [Array<String>] :ignore
       #   The advisories to ignore.
       #
+      # @option options [Array<String>] :filter
+      #   The criticalities to filter.
+      #
       # @yield [result]
       #   The given block will be passed the results of the scan.
       #
@@ -146,11 +149,16 @@ def scan_specs(options={})
         ignore = Set[]
         ignore += options[:ignore] if options[:ignore]
 
+        filter = Set[]
+        filter += options[:filter].map! { |current_filter| current_filter.downcase.to_sym } if options[:filter]
+
         @lockfile.specs.each do |gem|
           @database.check_gem(gem) do |advisory|
             is_ignored = ignore.intersect?(advisory.identifiers.to_set)
             next if is_ignored
 
+            next unless filter.empty? || filter.include?(advisory.criticality)
+
             yield UnpatchedGem.new(gem,advisory)
           end
         end
diff --git a/spec/integration_spec.rb b/spec/integration_spec.rb
index 0b31c5ce..4a8ef6e3 100644
--- a/spec/integration_spec.rb
+++ b/spec/integration_spec.rb
@@ -50,6 +50,24 @@
     end
   end
 
+  context "when auditing a bundle with filtered criticality" do
+    let(:bundle)    { 'unpatched_gems' }
+    let(:directory) { File.join('spec','bundle',bundle) }
+
+    let(:command) do
+      File.expand_path(File.join(File.dirname(__FILE__),'..','bin','bundler-audit -f medium'))
+    end
+
+    subject do
+      Dir.chdir(directory) { sh(command, :fail => true) }
+    end
+
+    it "should print advisories for filtered criticality" do
+      expect(subject).not_to include("Criticality: High")
+      expect(subject).to include("Criticality: Medium")
+    end
+  end
+
   context "when auditing a bundle with insecure sources" do
     let(:bundle)    { 'insecure_sources' }
     let(:directory) { File.join('spec','bundle',bundle) }
diff --git a/spec/scanner_spec.rb b/spec/scanner_spec.rb
index ce399b06..9fd6eba0 100644
--- a/spec/scanner_spec.rb
+++ b/spec/scanner_spec.rb
@@ -46,6 +46,16 @@
         expect(ids).not_to include('CVE-2013-0155')
       end
     end
+
+    context "when the :filter option is given" do
+      subject { scanner.scan(filter: ['high']) }
+
+      it "should return only filtered criticalities" do
+        criticalities = subject.map { |result| result.advisory.criticality }
+        expect(criticalities).not_to include(:medium, :low, nil)
+        expect(criticalities).to include(:high)
+      end
+    end
   end
 
   context "when auditing a bundle with insecure sources" do