feat: add functions to create and validate key

Two functions were added to create and validate a key during the agent registration process. Additionally, the --key argument is now optional.
wazuh · Oct 30, 2024 · 5f9f0ea · 5f9f0ea
1 parent bea5fb6
commit 5f9f0ea
Show file tree

Hide file tree

Showing 6 changed files with 85 additions and 21 deletions.
diff --git a/src/agent/agent_info/include/agent_info.hpp b/src/agent/agent_info/include/agent_info.hpp
@@ -15,11 +15,14 @@ class AgentInfo
     std::vector<std::string> GetGroups() const;
 
     void SetName(const std::string& name);
-    void SetKey(const std::string& key);
+    bool SetKey(const std::string& key);
     void SetUUID(const std::string& uuid);
     void SetGroups(const std::vector<std::string>& groupList);
 
 private:
+    std::string CreateKey();
+    bool ValidateKey(const std::string& key);
+
     std::string m_name;
     std::string m_key;
     std::string m_uuid;

diff --git a/src/agent/agent_info/src/agent_info.cpp b/src/agent/agent_info/src/agent_info.cpp
@@ -2,11 +2,19 @@
 
 #include <agent_info_persistance.hpp>
 
+#include <boost/random.hpp>
+#include <boost/random/random_device.hpp>
+#include <boost/random/uniform_int_distribution.hpp>
 #include <boost/uuid/uuid.hpp>
 #include <boost/uuid/uuid_generators.hpp>
 #include <boost/uuid/uuid_io.hpp>
 #include <utility>
 
+namespace
+{
+    constexpr size_t KEY_LENGTH = 32;
+}
+
 AgentInfo::AgentInfo()
 {
     AgentInfoPersistance agentInfoPersistance;
@@ -60,11 +68,25 @@ void AgentInfo::SetName(const std::string& name)
     m_name = name;
 }
 
-void AgentInfo::SetKey(const std::string& key)
+bool AgentInfo::SetKey(const std::string& key)
 {
     AgentInfoPersistance agentInfoPersistance;
-    agentInfoPersistance.SetKey(key);
-    m_key = key;
+    if (!key.empty())
+    {
+        if (!ValidateKey(key))
+        {
+            return false;
+        }
+        m_key = key;
+    }
+    else
+    {
+        m_key = CreateKey();
+    }
+
+    agentInfoPersistance.SetKey(m_key);
+
+    return true;
 }
 
 void AgentInfo::SetUUID(const std::string& uuid)
@@ -80,3 +102,31 @@ void AgentInfo::SetGroups(const std::vector<std::string>& groupList)
     agentInfoPersistance.SetGroups(groupList);
     m_groups = groupList;
 }
+
+std::string AgentInfo::CreateKey()
+{
+    constexpr std::string_view charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+
+    boost::random::mt19937 generator {boost::random::random_device {}()};
+    boost::random::uniform_int_distribution<size_t> distribution(0, charset.size() - 1);
+
+    std::string key;
+    for (size_t i = 0; i < KEY_LENGTH; ++i)
+    {
+        key += charset[distribution(generator)];
+    }
+
+    return key;
+}
+
+bool AgentInfo::ValidateKey(const std::string& key)
+{
+    if (key.length() != KEY_LENGTH)
+    {
+        return false;
+    }
+
+    return std::ranges::all_of(key, ::isalnum);
+
+    return true;
+}
diff --git a/src/agent/agent_info/tests/agent_info_test.cpp b/src/agent/agent_info/tests/agent_info_test.cpp
@@ -63,7 +63,7 @@ TEST_F(AgentInfoTest, TestSetName)
 TEST_F(AgentInfoTest, TestSetKey)
 {
     AgentInfo agentInfo;
-    const std::string newKey = "new_key";
+    const std::string newKey = "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj";
 
     agentInfo.SetKey(newKey);
     EXPECT_EQ(agentInfo.GetKey(), newKey);

diff --git a/src/agent/src/agent_registration.cpp b/src/agent/src/agent_registration.cpp
@@ -1,8 +1,8 @@
 #include <agent_registration.hpp>
 
-#include <logger.hpp>
-
 #include <boost/beast/http.hpp>
+#include <fmt/format.h>
+#include <iostream>
 #include <nlohmann/json.hpp>
 
 namespace http = boost::beast::http;
@@ -23,7 +23,11 @@ namespace agent_registration
         , m_password(std::move(password))
         , m_useHttps(!(m_configurationParser.GetConfig<std::string>("agent", "https_enabled") == "no"))
     {
-        m_agentInfo.SetKey(key);
+        if (!m_agentInfo.SetKey(key))
+        {
+            throw std::invalid_argument("--key argument must be alphanumeric and 32 characters in lengt");
+        }
+
         if (!name.empty())
         {
             m_agentInfo.SetName(name);
@@ -41,7 +45,7 @@ namespace agent_registration
 
         if (!token.has_value())
         {
-            LogError("Failed to authenticate with the manager");
+            std::cout << fmt::format("Failed to authenticate with the manager\n");
             return false;
         }
 
@@ -59,7 +63,7 @@ namespace agent_registration
 
         if (res.result() != http::status::ok)
         {
-            LogError("Registration error: {}.", res.result_int());
+            std::cout << fmt::format("Registration error: {}.\n", res.result_int());
             return false;
         }
 

diff --git a/src/agent/src/process_options.cpp b/src/agent/src/process_options.cpp
@@ -13,22 +13,29 @@ void RegisterAgent(const std::string& user,
                    const std::string& name,
                    const std::string& configFile)
 {
-    if (!user.empty() && !password.empty() && !key.empty())
+    if (!user.empty() && !password.empty())
     {
-        agent_registration::AgentRegistration reg(user, password, key, name, configFile);
-
-        http_client::HttpClient httpClient;
-        if (reg.Register(httpClient))
+        try
         {
-            std::cout << "wazuh-agent registered\n";
+            agent_registration::AgentRegistration reg(user, password, key, name, configFile);
+
+            http_client::HttpClient httpClient;
+            if (reg.Register(httpClient))
+            {
+                std::cout << "wazuh-agent registered\n";
+            }
+            else
+            {
+                std::cout << "wazuh-agent registration failed\n";
+            }
         }
-        else
+        catch (const std::exception& e)
         {
-            std::cout << "wazuh-agent registration failed\n";
+            std::cerr << "Error: " << e.what() << '\n';
         }
     }
     else
     {
-        std::cout << fmt::format("{}, {}, and {} args are mandatory\n", OPT_USER, OPT_PASSWORD, OPT_KEY);
+        std::cout << fmt::format("{} and {} args are mandatory\n", OPT_USER, OPT_PASSWORD);
     }
 }
diff --git a/src/agent/tests/agent_registration_test.cpp b/src/agent/tests/agent_registration_test.cpp
@@ -66,9 +66,9 @@ class RegisterTest : public ::testing::Test
 protected:
     void SetUp() override
     {
-        agent = std::make_unique<AgentInfo>("agent_name", "agent_key", "agent_uuid");
+        agent = std::make_unique<AgentInfo>("agent_name", "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj", "agent_uuid");
         registration = std::make_unique<agent_registration::AgentRegistration>(
-            "user", "password", "agent_key", "agent_name", std::nullopt);
+            "user", "password", "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj", "agent_name", std::nullopt);
     }
 
     std::unique_ptr<AgentInfo> agent;