Skip to content

Commit

Permalink
Introduce clean-ups for actions in the templated language.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 726097872
Change-Id: Ifa44330130e06acb5029c221803cb59460f981ce
  • Loading branch information
tooryx authored and copybara-github committed Feb 12, 2025
1 parent 6e1a001 commit 6fb72e1
Show file tree
Hide file tree
Showing 3 changed files with 151 additions and 20 deletions.
12 changes: 9 additions & 3 deletions templated/templateddetector/proto/templated_plugin.proto
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,17 @@ message PluginAction {
// workflows. It must be named using the `[a-zA-Z0-9_]` character set.
string name = 1;

// A set of cleanup action to be executed if this action is successful.
// Once the current action succeed, the cleanups are registered and will
// always be executed after the last workflow action, whether it is successful
// or not.
repeated string cleanup_actions = 2;

// Each action can have one of the following types.
oneof any_action {
HttpAction http_request = 2;
CallbackServerAction callback_server = 3;
UtilityAction utility = 4;
HttpAction http_request = 3;
CallbackServerAction callback_server = 4;
UtilityAction utility = 5;
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -147,29 +147,63 @@ private final boolean runWorkflowForService(NetworkService service, PluginWorkfl
environment.set(parameter.getName(), value);
}

// Run every individual action and inventorize clean up actions.
boolean success = true;
ImmutableList.Builder<String> cleanupActionsBuilder = ImmutableList.builder();
for (String actionName : workflow.getActionsList()) {
PluginAction action = this.actionsCache.get(actionName);
if (!dispatchAction(service, action, environment)) {
logger.atInfo().log("No vulnerability found because action '%s' failed.", actionName);
return false;
success = false;
break;
}

cleanupActionsBuilder.addAll(action.getCleanupActionsList());
}

var cleanupActions = cleanupActionsBuilder.build();
if (cleanupActions.isEmpty()) {
return success;
}

// Run all the clean up actions that were registered.
logger.atInfo().log("Running %d cleanup action(s)", cleanupActions.size());
for (String cleanupActionName : cleanupActions) {
PluginAction cleanupAction = this.actionsCache.get(cleanupActionName);
if (!dispatchAction(service, cleanupAction, environment)) {
logger.atWarning().log(
"Cleanup action '%s' failed: manual clean up of service on port %d might be needed",
cleanupActionName, service.getNetworkEndpoint().getPort().getPortNumber());
}
}

return true;
return success;
}

private final DetectionReportList useWorkflow(
TargetInfo targetInfo,
ImmutableList<NetworkService> matchedService,
PluginWorkflow workflow) {
// First we precheck that all registered actions exists.
// First we check that all registered actions exists.
for (String actionName : workflow.getActionsList()) {
if (!this.actionsCache.containsKey(actionName)) {
throw new IllegalArgumentException(
String.format("Plugin definition error: action '%s' not found.", actionName));
}
}

// Also check that all registered cleanups are valid.
for (PluginAction action : this.actionsCache.values()) {
for (String cleanupActionName : action.getCleanupActionsList()) {
if (!this.actionsCache.containsKey(cleanupActionName)) {
throw new IllegalArgumentException(
String.format(
"Plugin definition error: cleanup action '%s' defined in action '%s' not found.",
cleanupActionName, action.getName()));
}
}
}

return DetectionReportList.newBuilder()
.addAllDetectionReports(
matchedService.stream()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,25 +51,44 @@ public void nextBytes(byte[] bytes) {
Arrays.fill(bytes, (byte) 0xFF);
}
};
private static final PluginAction ACTION_RETURNS_TRUE =
PluginAction.newBuilder()
.setName("action_returns_true")
.setHttpRequest(
HttpAction.newBuilder()
.setMethod(HttpAction.HttpMethod.GET)
.addUri("/OK")
.setResponse(HttpAction.HttpResponse.newBuilder().setHttpStatus(200)))
.build();
private static final PluginAction ACTION_RETURNS_FALSE =
PluginAction.newBuilder()
.setName("action_returns_false")
.setHttpRequest(
HttpAction.newBuilder()
.setMethod(HttpAction.HttpMethod.GET)
.addUri("/NOTFOUND")
.setResponse(HttpAction.HttpResponse.newBuilder().setHttpStatus(200)))
.build();
private static final PluginAction ACTION_CLEANUP =
ACTION_RETURNS_TRUE.toBuilder()
.setName("action_cleanup")
.setHttpRequest(
HttpAction.newBuilder().setMethod(HttpAction.HttpMethod.GET).addUri("/CLEANUP"))
.build();
private static final TemplatedPlugin BASE_PROTO =
TemplatedPlugin.newBuilder()
.setInfo(PluginInfo.newBuilder().setName("ExampleTemplated"))
.addActions(ACTION_CLEANUP)
.addActions(ACTION_RETURNS_TRUE)
.addActions(
PluginAction.newBuilder()
.setName("action_returns_true")
.setHttpRequest(
HttpAction.newBuilder()
.setMethod(HttpAction.HttpMethod.GET)
.addUri("/OK")
.setResponse(HttpAction.HttpResponse.newBuilder().setHttpStatus(200))))
ACTION_RETURNS_TRUE.toBuilder()
.setName("action_returns_true_with_cleanup")
.addCleanupActions("action_cleanup"))
.addActions(ACTION_RETURNS_FALSE)
.addActions(
PluginAction.newBuilder()
.setName("action_returns_false")
.setHttpRequest(
HttpAction.newBuilder()
.setMethod(HttpAction.HttpMethod.GET)
.addUri("/NOTFOUND")
.setResponse(HttpAction.HttpResponse.newBuilder().setHttpStatus(200))))
ACTION_RETURNS_FALSE.toBuilder()
.setName("action_returns_false_with_cleanup")
.addCleanupActions("action_cleanup"))
.build();

private MockWebServer mockWebServer;
Expand Down Expand Up @@ -214,6 +233,22 @@ public void detect_unknownActionNameInWorkflow_throwsException() {
assertThat(this.mockWebServer.getRequestCount()).isEqualTo(0);
}

@Test
public void detect_unknownCleanupNameInAction_throwsException() {
var proto =
BASE_PROTO.toBuilder()
.addActions(
PluginAction.newBuilder()
.setName("action_with_undefined_cleanup")
.addCleanupActions("undefined_cleanup"))
.addWorkflows(PluginWorkflow.newBuilder().addActions("action_with_undefined_cleanup"))
.build();
var detector = setupDetector(proto);

assertThrows(
IllegalArgumentException.class, () -> detector.detect(this.targetInfo, this.httpServices));
}

@Test
public void detect_variableFromWorkflow_propagatedAndReturnsFindings()
throws InterruptedException {
Expand Down Expand Up @@ -277,6 +312,62 @@ public void detect_customFinding_returnsCustomFinding() throws InterruptedExcept
.containsExactly(expect);
}

@Test
public void detect_cleanupOnFailingAction_cleanupDoesNotRun() throws InterruptedException {
var proto =
BASE_PROTO.toBuilder()
.addWorkflows(
PluginWorkflow.newBuilder().addActions("action_returns_false_with_cleanup"))
.build();
var detector = setupDetector(proto);

assertThat(detector.detect(this.targetInfo, this.httpServices).getDetectionReportsCount())
.isEqualTo(0);
assertThat(this.mockWebServer.getRequestCount()).isEqualTo(1);
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/NOTFOUND");
}

@Test
public void detect_cleanupOnSuccessfulAction_cleanupRunsAfterLastSuccessAction()
throws InterruptedException {
var proto =
BASE_PROTO.toBuilder()
.addWorkflows(
PluginWorkflow.newBuilder()
.addActions("action_returns_true_with_cleanup")
.addActions("action_returns_true"))
.build();
var detector = setupDetector(proto);

assertThat(detector.detect(this.targetInfo, this.httpServices).getDetectionReportsCount())
.isEqualTo(1);
assertThat(this.mockWebServer.getRequestCount()).isEqualTo(3);
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/OK");
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/OK");
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/CLEANUP");
}

@Test
public void detect_cleanupOnSuccessfulAction_cleanupRunsAfterLastFailAction()
throws InterruptedException {
var proto =
BASE_PROTO.toBuilder()
.addWorkflows(
PluginWorkflow.newBuilder()
.addActions("action_returns_true_with_cleanup")
.addActions("action_returns_false")
.addActions("action_returns_true"))
.build();
var detector = setupDetector(proto);

assertThat(detector.detect(this.targetInfo, this.httpServices).getDetectionReportsCount())
.isEqualTo(0);
assertThat(this.mockWebServer.getRequestCount()).isEqualTo(3);
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/OK");
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/NOTFOUND");
assertThat(this.mockWebServer.takeRequest().getPath()).isEqualTo("/CLEANUP");
}

private TemplatedDetector setupDetector(TemplatedPlugin proto) {
TemplatedDetector detector = new TemplatedDetector(proto);
Guice.createInjector(
Expand Down

0 comments on commit 6fb72e1

Please sign in to comment.