Skip to content

Commit

Permalink
Clean up and refactor orders code
Browse files Browse the repository at this point in the history
  • Loading branch information
@f-galland
f-galland committed Dec 19, 2024
1 parent dbbdaf3 commit a2f965e
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 79 deletions.
37 changes: 2 additions & 35 deletions ...ins/command-manager/src/main/java/com/wazuh/commandmanager/jobscheduler/SearchThread.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,6 @@
import org.opensearch.client.Client;
import org.opensearch.common.action.ActionFuture;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.common.xcontent.XContentHelper;
import org.opensearch.common.xcontent.XContentType;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.*;
Expand All @@ -42,7 +39,6 @@
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.sort.SortOrder;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.security.AccessController;
Expand All @@ -53,10 +49,7 @@
import java.util.concurrent.ExecutionException;

import com.wazuh.commandmanager.CommandManagerPlugin;
import com.wazuh.commandmanager.model.Command;
import com.wazuh.commandmanager.model.Document;
import com.wazuh.commandmanager.model.Order;
import com.wazuh.commandmanager.model.Status;
import com.wazuh.commandmanager.model.*;
import com.wazuh.commandmanager.settings.PluginSettings;
import com.wazuh.commandmanager.utils.httpclient.AuthHttpRestClient;

Expand Down Expand Up @@ -121,33 +114,7 @@ public static <T> T getNestedObject(Map<String, Object> map, String key, Class<T
@SuppressWarnings("unchecked")
public void handlePage(SearchResponse searchResponse) throws IllegalStateException {
SearchHits searchHits = searchResponse.getHits();
String payload = null;
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
// Start an XContentBuilder array named "orders"
builder.startObject();
builder.startArray(Order.ORDERS);
// Iterate over search results
for (SearchHit hit : searchHits) {
// Create a parser for each SearchHit
XContentParser parser =
XContentHelper.createParser(
NamedXContentRegistry.EMPTY,
DeprecationHandler.IGNORE_DEPRECATIONS,
hit.getSourceRef(),
XContentType.JSON);
// Parse the hit's order
Order order = Order.parseSearchHit(hit);
// Add the current order to the XContentBuilder array
assert order != null;
order.toXContent(builder, ToXContent.EMPTY_PARAMS);
}
// Close the object and prepare it for delivery
builder.endArray();
builder.endObject();
payload = builder.toString();
} catch (IOException e) {
log.error("Error building payload from hit: {}", e.getMessage());
}
String payload = Orders.getOrders(searchHits);
final SimpleHttpResponse response = deliverOrders(payload);
if (response == null) {
log.error("No reply from server.");
Expand Down
57 changes: 28 additions & 29 deletions plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Command.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,31 +97,30 @@ public static Command parse(XContentParser parser)
Action action = null;

while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
if (!parser.currentToken().equals(XContentParser.Token.FIELD_NAME)) {
continue;
}
String fieldName = parser.currentName();

parser.nextToken();
switch (fieldName) {
case SOURCE:
source = parser.text();
break;
case Target.TARGET:
target = Target.parse(parser);
break;
case TIMEOUT:
timeout = parser.intValue();
break;
case USER:
user = parser.text();
break;
case Action.ACTION:
action = Action.parse(parser);
break;
default:
parser.skipChildren();
break;
if (parser.currentToken().equals(XContentParser.Token.FIELD_NAME)) {
String fieldName = parser.currentName();

parser.nextToken();
switch (fieldName) {
case SOURCE:
source = parser.text();
break;
case Target.TARGET:
target = Target.parse(parser);
break;
case TIMEOUT:
timeout = parser.intValue();
break;
case USER:
user = parser.text();
break;
case Action.ACTION:
action = Action.parse(parser);
break;
default:
parser.skipChildren();
break;
}
}
}

Expand Down Expand Up @@ -183,19 +182,19 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
}

public Action getAction() {
return action;
return this.action;
}

public String getSource() {
return source;
return this.source;
}

public Target getTarget() {
return target;
return this.target;
}

public String getUser() {
return user;
return this.user;
}

@Override
Expand Down
33 changes: 18 additions & 15 deletions plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Order.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import org.opensearch.search.SearchHit;

import java.io.IOException;
import java.util.Objects;

import reactor.util.annotation.NonNull;

Expand Down Expand Up @@ -84,27 +85,29 @@ public static Order parseSearchHit(SearchHit hit) {
// which corresponds to end of data
while (parser.nextToken() != null) {
// Look for FIELD_NAME JsonToken s
if (!parser.currentToken().equals(XContentParser.Token.FIELD_NAME)) {
continue;
}
String fieldName = parser.currentName();
if (fieldName.equals(Command.COMMAND)) {
// Parse Command
command = Command.parse(parser);
} else {
parser.skipChildren();
if (parser.currentToken().equals(XContentParser.Token.FIELD_NAME)) {
String fieldName = parser.currentName();
if (fieldName.equals(Command.COMMAND)) {
// Parse Command
command = Command.parse(parser);
} else {
parser.skipChildren();
}
}
}
// Create a new Order object with the Command's fields
assert command != null;
return new Order(
command.getSource(),
command.getTarget(),
command.getUser(),
command.getAction(),
hit.getId());
Objects.requireNonNull(command).getSource(),
Objects.requireNonNull(command).getTarget(),
Objects.requireNonNull(command).getUser(),
Objects.requireNonNull(command).getAction(),
Objects.requireNonNull(hit).getId());
} catch (IOException e) {
log.error("Order could not be parsed: {}", e.getMessage());
} catch (NullPointerException e) {
log.error(
"Could not create Order object. One or more of the constructor's arguments was null: {}",
e.getMessage());
}
return null;
}
Expand Down
58 changes: 58 additions & 0 deletions plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Orders.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
/*
* Copyright (C) 2024, Wazuh Inc.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
package com.wazuh.commandmanager.model;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.search.SearchHit;
import org.opensearch.search.SearchHits;

import java.io.IOException;
import java.util.Objects;

public class Orders {

private static final Logger log = LogManager.getLogger(Orders.class);

public static String getOrders(SearchHits searchHits) {
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
// Start an XContentBuilder array named "orders"
builder.startObject();
builder.startArray(Order.ORDERS);
// Iterate over search results
for (SearchHit hit : searchHits) {
// Parse the hit's order
Order order = Order.parseSearchHit(hit);
// Add the current order to the XContentBuilder array
Objects.requireNonNull(order).toXContent(builder, ToXContent.EMPTY_PARAMS);
}
// Close the object and prepare it for delivery
builder.endArray();
builder.endObject();
return builder.toString();
} catch (IOException e) {
log.error("Error building payload from hit: {}", e.getMessage());
} catch (NullPointerException e) {
log.error(
"Exception found when building order payload. Null Order: {}", e.getMessage());
}
return null;
}
}

0 comments on commit a2f965e

Please sign in to comment.