Introduce the RequiresCallbackServer annotation that ensures a plug…

…in only gets executed if the callback server is present.

PiperOrigin-RevId: 667556868
Change-Id: I8641612e85372a89812e2173705da999e640c35e

main/src/test/java/com/google/tsunami/main/cli/TsunamiCliTest.java

@@ -27,8 +27,10 @@
 import com.google.tsunami.common.config.ConfigModule;
 import com.google.tsunami.common.config.TsunamiConfig;
 import com.google.tsunami.common.data.NetworkEndpointUtils;
+import com.google.tsunami.common.net.http.HttpClientModule;
 import com.google.tsunami.common.time.testing.FakeUtcClockModule;
 import com.google.tsunami.main.cli.server.RemoteServerLoaderModule;
+import com.google.tsunami.plugin.payload.PayloadGeneratorModule;
 import com.google.tsunami.plugin.testing.FailedVulnDetectorBootstrapModule;
 import com.google.tsunami.plugin.testing.FakePluginExecutionModule;
 import com.google.tsunami.plugin.testing.FakePortScanner;
@@ -62,6 +64,7 @@
 import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.net.URL;
+import java.security.SecureRandom;
 import java.util.concurrent.ExecutionException;
 import java.util.stream.Stream;
 import javax.inject.Inject;
@@ -98,6 +101,8 @@ private boolean runCli(ImmutableMap<String, Object> rawConfigData, String... arg
                 @Override
                 protected void configure() {
                   bind(ScanResultsArchiver.class).toInstance(scanResultsArchiver);
+                  install(new HttpClientModule.Builder().build());
+                  install(new PayloadGeneratorModule(new SecureRandom()));
                   install(new ConfigModule(scanResult, TsunamiConfig.fromYamlData(rawConfigData)));
                   install(new CliOptionsModule(scanResult, "TsunamiCliTest", args));
                   install(new FakeUtcClockModule());
@@ -271,6 +276,8 @@ public void run_whenScanFailed_generatesFailedScanResults()
                 @Override
                 protected void configure() {
                   bind(ScanResultsArchiver.class).toInstance(scanResultsArchiver);
+                  install(new HttpClientModule.Builder().build());
+                  install(new PayloadGeneratorModule(new SecureRandom()));
                   install(
                       new ConfigModule(scanResult, TsunamiConfig.fromYamlData(ImmutableMap.of())));
                   install(

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

@@ -25,6 +25,7 @@
 import com.google.tsunami.plugin.annotations.ForSoftware;
 import com.google.tsunami.plugin.annotations.ForWebService;
 import com.google.tsunami.plugin.annotations.PluginInfo;
+import com.google.tsunami.plugin.annotations.RequiresCallbackServer;
 import java.util.Optional;
 
 /** A data class that captures all the definition details about a {@link TsunamiPlugin}. */
@@ -40,6 +41,8 @@ abstract class PluginDefinition {
 
   abstract Optional<ForOperatingSystemClass> targetOperatingSystemClass();
 
+  abstract boolean requiresCallbackServer();
+
   /**
    * Unique identifier for the plugin.
    *
@@ -82,6 +85,7 @@ public static PluginDefinition forPlugin(Class<? extends TsunamiPlugin> pluginCl
     boolean isForWebService = pluginClazz.isAnnotationPresent(ForWebService.class);
     Optional<ForOperatingSystemClass> targetOperatingSystemClass =
         Optional.ofNullable(pluginClazz.getAnnotation(ForOperatingSystemClass.class));
+    boolean requiresCallbackServer = pluginClazz.isAnnotationPresent(RequiresCallbackServer.class);
 
     checkState(
         pluginInfo.isPresent(),
@@ -93,7 +97,8 @@ public static PluginDefinition forPlugin(Class<? extends TsunamiPlugin> pluginCl
         targetServiceName,
         targetSoftware,
         isForWebService,
-        targetOperatingSystemClass);
+        targetOperatingSystemClass,
+        requiresCallbackServer);
   }
 
   /**
@@ -109,6 +114,7 @@ public static PluginDefinition forRemotePlugin(PluginInfo remotePluginInfo) {
         Optional.empty(),
         Optional.empty(),
         false,
-        Optional.empty());
+        Optional.empty(),
+        false);
   }
 }

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

@@ -15,6 +15,7 @@
  */
 package com.google.tsunami.plugin;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.tsunami.common.data.NetworkServiceUtils.isWebService;
 import static java.util.Arrays.stream;
@@ -41,10 +42,13 @@
  */
 public class PluginManager {
   private final Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins;
+  private final TcsClient tcsClient;
 
   @Inject
-  PluginManager(Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins) {
+  PluginManager(
+      Map<PluginDefinition, Provider<TsunamiPlugin>> tsunamiPlugins, TcsClient tcsClient) {
     this.tsunamiPlugins = tsunamiPlugins;
+    this.tcsClient = checkNotNull(tcsClient);
   }
 
   /**
@@ -100,6 +104,7 @@ public ImmutableList<PluginMatchingResult<VulnDetector>> getVulnDetectors(
       ReconnaissanceReport reconnaissanceReport) {
     return tsunamiPlugins.entrySet().stream()
         .filter(entry -> isVulnDetector(entry.getKey()))
+        .filter(entry -> matchCurrentCallbackServerSetup(entry.getKey()))
         .map(entry -> matchAllVulnDetectors(entry.getKey(), entry.getValue(), reconnaissanceReport))
         .flatMap(Streams::stream)
         .collect(toImmutableList());
@@ -110,6 +115,14 @@ private static boolean isVulnDetector(PluginDefinition pluginDefinition) {
         || pluginDefinition.type().equals(PluginType.REMOTE_VULN_DETECTION);
   }
 
+  private boolean matchCurrentCallbackServerSetup(PluginDefinition pluginDefinition) {
+    if (tcsClient.isCallbackServerEnabled()) {
+      return true;
+    }
+
+    return !pluginDefinition.requiresCallbackServer();
+  }
+
   private static Optional<PluginMatchingResult<VulnDetector>> matchAllVulnDetectors(
       PluginDefinition pluginDefinition,
       Provider<TsunamiPlugin> vulnDetectorProvider,
@@ -310,13 +323,16 @@ public static <T extends TsunamiPlugin> Builder<T> builder() {
     @AutoValue.Builder
     public abstract static class Builder<T extends TsunamiPlugin> {
       public abstract Builder<T> setPluginDefinition(PluginDefinition value);
+
       public abstract Builder<T> setTsunamiPlugin(T value);
 
       abstract ImmutableList.Builder<NetworkService> matchedServicesBuilder();
+
       public Builder<T> addMatchedService(NetworkService networkService) {
         matchedServicesBuilder().add(networkService);
         return this;
       }
+
       public Builder<T> addAllMatchedServices(Iterable<NetworkService> networkServices) {
         matchedServicesBuilder().addAll(networkServices);
         return this;

plugin/src/main/java/com/google/tsunami/plugin/annotations/RequiresCallbackServer.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024 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.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation for marking that a plugin can only run if the callback server is enabled.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+public @interface RequiresCallbackServer {}