New command line flags --detectors-include and `--detectors-exclude…

…` to fine tune which scanners to execute during the scan.

PiperOrigin-RevId: 722986341
Change-Id: I89234e177db000ba6763717132c01fbb37cc84da

plugin/src/main/java/com/google/tsunami/plugin/PluginManager.java

@@ -17,12 +17,14 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.tsunami.common.data.NetworkServiceUtils.isWebService;
 import static java.util.Arrays.stream;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.tsunami.plugin.annotations.ForOperatingSystemClass;
 import com.google.tsunami.proto.MatchedPlugin;
@@ -43,12 +45,26 @@
 public class PluginManager {
   private final Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins;
   private final TcsClient tcsClient;
+  private final ImmutableSet<String> detectorsInclude;
+  private final ImmutableSet<String> detectorsExclude;
 
   @Inject
   PluginManager(
-      Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins, TcsClient tcsClient) {
+      Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins,
+      TcsClient tcsClient,
+      PluginManagerCliOptions pluginManagerCliOptions) {
     this.tsunamiPlugins = tsunamiPlugins;
     this.tcsClient = checkNotNull(tcsClient);
+    detectorsInclude = getDetectorNames(pluginManagerCliOptions.detectorsInclude);
+    detectorsExclude = getDetectorNames(pluginManagerCliOptions.detectorsExclude);
+  }
+
+  private static ImmutableSet<String> getDetectorNames(String detectorNames) {
+    if (detectorNames == null) {
+      return ImmutableSet.of();
+    } else {
+      return stream(detectorNames.split(",")).map(String::trim).collect(toImmutableSet());
+    }
   }
 
   /**
@@ -105,11 +121,25 @@ public ImmutableList<PluginMatchingResult<VulnDetector>> getVulnDetectors(
     return tsunamiPlugins.entrySet().stream()
         .filter(entry -> isVulnDetector(entry.getKey()))
         .filter(entry -> matchCurrentCallbackServerSetup(entry.getKey()))
+        .filter(entry -> filterPluginByCliOptions(entry.getKey()))
         .map(entry -> matchAllVulnDetectors(entry.getKey(), entry.getValue(), reconnaissanceReport))
         .flatMap(Streams::stream)
         .collect(toImmutableList());
   }
 
+  private static boolean isPluginListed(
+      PluginDefinition pluginDefinition, ImmutableSet<String> pluginNames, boolean defaultValue) {
+    if (pluginNames.isEmpty()) {
+      return defaultValue;
+    }
+    return pluginNames.contains(pluginDefinition.name());
+  }
+
+  private boolean filterPluginByCliOptions(PluginDefinition pluginDefinition) {
+    return isPluginListed(pluginDefinition, detectorsInclude, true)
+        && !isPluginListed(pluginDefinition, detectorsExclude, false);
+  }
+
   private static boolean isVulnDetector(PluginDefinition pluginDefinition) {
     return pluginDefinition.type().equals(PluginType.VULN_DETECTION)
         || pluginDefinition.type().equals(PluginType.REMOTE_VULN_DETECTION);

plugin/src/main/java/com/google/tsunami/plugin/PluginManagerCliOptions.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.tsunami.plugin;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.tsunami.common.cli.CliOption;
+
+/**
+ * Command line arguments for the PluginManager.
+ *
+ * <p>Which detectors to include/exclude? Matching is executed against the name of the detector
+ * (e.g. RsyncRceDetector).
+ *
+ * <p>To execute one single detector, override `detectors-include`, e.g.:
+ * `--detectors-include=RsyncRceDetector`.
+ *
+ * <p>To disable a single detector, override `detectors-exclude`, e.g.
+ * `--detectors-exclude=RsyncRceDetector`.
+ */
+@Parameters(separators = "=")
+public final class PluginManagerCliOptions implements CliOption {
+  @Parameter(
+      names = "--detectors-include",
+      description =
+          "Comma separated list of detector names to include in the scan. By default, all detectors"
+              + " are included.")
+  public String detectorsInclude;
+
+  @Parameter(
+      names = "--detectors-exclude",
+      description =
+          "Comma separated list of detector names to exclude from the scan. By default no detectors"
+              + " are skipped.")
+  public String detectorsExclude;
+
+  // Validations are done in {@link PayloadGeneratorModule}.
+  @Override
+  public void validate() {}
+}

plugin/src/test/java/com/google/tsunami/plugin/PluginManagerTest.java

@@ -23,6 +23,7 @@
 import com.google.inject.AbstractModule;
 import com.google.inject.Guice;
 import com.google.inject.multibindings.MapBinder;
+import com.google.tsunami.common.cli.CliOptionsModule;
 import com.google.tsunami.common.data.NetworkEndpointUtils;
 import com.google.tsunami.common.net.http.HttpClientModule;
 import com.google.tsunami.plugin.PluginManager.PluginMatchingResult;
@@ -54,6 +55,8 @@
 import com.google.tsunami.proto.TargetServiceName;
 import com.google.tsunami.proto.TargetSoftware;
 import com.google.tsunami.proto.TransportProtocol;
+import io.github.classgraph.ClassGraph;
+import io.github.classgraph.ScanResult;
 import java.security.SecureRandom;
 import java.util.List;
 import java.util.Optional;
@@ -934,6 +937,100 @@ public void getVulnDetectors_whenRemoteDetectorWithServiceNameHasNoMatch_returns
     }
   }
 
+  @Test
+  public void getVulnDetectors_whenDetectorsIncludeIsOverridden_returnsMatchingVulnDetector() {
+    NetworkService fakeNetworkService1 =
+        NetworkService.newBuilder()
+            .setNetworkEndpoint(NetworkEndpointUtils.forIpAndPort("1.1.1.1", 80))
+            .setTransportProtocol(TransportProtocol.TCP)
+            .setServiceName("http")
+            .build();
+    NetworkService fakeNetworkService2 =
+        NetworkService.newBuilder()
+            .setNetworkEndpoint(NetworkEndpointUtils.forIpAndPort("1.1.1.1", 443))
+            .setTransportProtocol(TransportProtocol.TCP)
+            .setServiceName("https")
+            .build();
+    ReconnaissanceReport fakeReconnaissanceReport =
+        ReconnaissanceReport.newBuilder()
+            .setTargetInfo(TargetInfo.getDefaultInstance())
+            .addNetworkServices(fakeNetworkService1)
+            .addNetworkServices(fakeNetworkService2)
+            .build();
+    try (ScanResult scanResult = new ClassGraph().enableAllInfo().scan()) {
+      PluginManager pluginManager =
+          Guice.createInjector(
+                  new CliOptionsModule(
+                      scanResult,
+                      "TsunamiCliTest",
+                      new String[] {"--detectors-include=Blabla1, FakeVulnDetector, Blabla2"}),
+                  new HttpClientModule.Builder().build(),
+                  new PayloadGeneratorModule(new SecureRandom()),
+                  new FakePortScannerBootstrapModule(),
+                  new FakeServiceFingerprinterBootstrapModule(),
+                  new FakeVulnDetectorBootstrapModule(),
+                  new FakeVulnDetectorBootstrapModule2())
+              .getInstance(PluginManager.class);
+
+      ImmutableList<PluginMatchingResult<VulnDetector>> vulnDetectors =
+          pluginManager.getVulnDetectors(fakeReconnaissanceReport);
+
+      assertThat(
+              vulnDetectors.stream()
+                  .map(pluginMatchingResult -> pluginMatchingResult.tsunamiPlugin().getClass()))
+          .containsExactly(FakeVulnDetector.class);
+      assertThat(vulnDetectors.stream().map(PluginMatchingResult::matchedServices))
+          .containsExactly(fakeReconnaissanceReport.getNetworkServicesList());
+    }
+  }
+
+  @Test
+  public void getVulnDetectors_whenDetectorsExcludeIsOverridden_returnsMatchingVulnDetector() {
+    NetworkService fakeNetworkService1 =
+        NetworkService.newBuilder()
+            .setNetworkEndpoint(NetworkEndpointUtils.forIpAndPort("1.1.1.1", 80))
+            .setTransportProtocol(TransportProtocol.TCP)
+            .setServiceName("http")
+            .build();
+    NetworkService fakeNetworkService2 =
+        NetworkService.newBuilder()
+            .setNetworkEndpoint(NetworkEndpointUtils.forIpAndPort("1.1.1.1", 443))
+            .setTransportProtocol(TransportProtocol.TCP)
+            .setServiceName("https")
+            .build();
+    ReconnaissanceReport fakeReconnaissanceReport =
+        ReconnaissanceReport.newBuilder()
+            .setTargetInfo(TargetInfo.getDefaultInstance())
+            .addNetworkServices(fakeNetworkService1)
+            .addNetworkServices(fakeNetworkService2)
+            .build();
+    try (ScanResult scanResult = new ClassGraph().enableAllInfo().scan()) {
+      PluginManager pluginManager =
+          Guice.createInjector(
+                  new CliOptionsModule(
+                      scanResult,
+                      "TsunamiCliTest",
+                      new String[] {"--detectors-exclude=FakeVulnDetector"}),
+                  new HttpClientModule.Builder().build(),
+                  new PayloadGeneratorModule(new SecureRandom()),
+                  new FakePortScannerBootstrapModule(),
+                  new FakeServiceFingerprinterBootstrapModule(),
+                  new FakeVulnDetectorBootstrapModule(),
+                  new FakeVulnDetectorBootstrapModule2())
+              .getInstance(PluginManager.class);
+
+      ImmutableList<PluginMatchingResult<VulnDetector>> vulnDetectors =
+          pluginManager.getVulnDetectors(fakeReconnaissanceReport);
+
+      assertThat(
+              vulnDetectors.stream()
+                  .map(pluginMatchingResult -> pluginMatchingResult.tsunamiPlugin().getClass()))
+          .containsExactly(FakeVulnDetector2.class);
+      assertThat(vulnDetectors.stream().map(PluginMatchingResult::matchedServices))
+          .containsExactly(fakeReconnaissanceReport.getNetworkServicesList());
+    }
+  }
+
   @PluginInfo(
       type = PluginType.SERVICE_FINGERPRINT,
       name = "NoAnnotationFingerprinter",