From a2f965e6f6a543092609935d3ee6b2e19241eed5 Mon Sep 17 00:00:00 2001
From: f-galland <federico.galland@wazuh.com>
Date: Thu, 19 Dec 2024 13:49:32 -0300
Subject: [PATCH] Clean up and refactor orders code

---
 .../jobscheduler/SearchThread.java            | 37 +-----------
 .../wazuh/commandmanager/model/Command.java   | 57 +++++++++---------
 .../com/wazuh/commandmanager/model/Order.java | 33 ++++++-----
 .../wazuh/commandmanager/model/Orders.java    | 58 +++++++++++++++++++
 4 files changed, 106 insertions(+), 79 deletions(-)
 create mode 100644 plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Orders.java

diff --git a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/jobscheduler/SearchThread.java b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/jobscheduler/SearchThread.java
index 462b5f3..d9a773d 100644
--- a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/jobscheduler/SearchThread.java
+++ b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/jobscheduler/SearchThread.java
@@ -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.*;
@@ -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;
@@ -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;
 
@@ -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.");
diff --git a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Command.java b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Command.java
index e6dbf90..350751c 100644
--- a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Command.java
+++ b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Command.java
@@ -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;
+                }
             }
         }
 
@@ -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
diff --git a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Order.java b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Order.java
index 161a860..b75b09d 100644
--- a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Order.java
+++ b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Order.java
@@ -24,6 +24,7 @@
 import org.opensearch.search.SearchHit;
 
 import java.io.IOException;
+import java.util.Objects;
 
 import reactor.util.annotation.NonNull;
 
@@ -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;
     }
diff --git a/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Orders.java b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Orders.java
new file mode 100644
index 0000000..a2b6b5d
--- /dev/null
+++ b/plugins/command-manager/src/main/java/com/wazuh/commandmanager/model/Orders.java
@@ -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;
+    }
+}