Merge pull request #319 from wazuh/fix/314-registration-failure

TomasTurina · web-flow · commit a2405d2ff442 · 2024-11-19T16:10:36.000-03:00
Fix Agent Registration Behavior
diff --git a/src/agent/agent_info/include/agent_info.hpp b/src/agent/agent_info/include/agent_info.hpp
@@ -41,20 +41,20 @@ class AgentInfo
     /// @return A vector of the agent's groups.
     std::vector<std::string> GetGroups() const;
 
-    /// @brief Sets the agent's name.
+    /// @brief Sets the agent's name. The change is not saved to the database until `Save` is called.
     /// @param name The agent's new name.
     void SetName(const std::string& name);
 
-    /// @brief Sets the agent's key.
+    /// @brief Sets the agent's key. The change is not saved to the database until `Save` is called.
     /// @param key The agent's new key.
     /// @return True if the key was successfully set, false otherwise.
     bool SetKey(const std::string& key);
 
-    /// @brief Sets the agent's UUID.
+    /// @brief Sets the agent's UUID. The change is not saved to the database until `Save` is called.
     /// @param uuid The agent's new UUID.
     void SetUUID(const std::string& uuid);
 
-    /// @brief Sets the agent's groups.
+    /// @brief Sets the agent's groups. The change is not saved to the database until `Save` is called.
     /// @param groupList A vector of the agent's new groups.
     void SetGroups(const std::vector<std::string>& groupList);
 
@@ -75,6 +75,9 @@ class AgentInfo
     /// @return A string with all information about the agent.
     std::string GetMetadataInfo(const bool agentIsRegistering) const;
 
+    /// @brief Saves the agent's information to the database.
+    void Save() const;
+
 private:
     /// @brief Creates a random key for the agent.
     ///
diff --git a/src/agent/agent_info/src/agent_info.cpp b/src/agent/agent_info/src/agent_info.cpp
@@ -27,7 +27,6 @@ AgentInfo::AgentInfo(std::function<nlohmann::json()> getOSInfo, std::function<nl
     if (m_uuid.empty())
     {
         m_uuid = boost::uuids::to_string(boost::uuids::random_generator()());
-        agentInfoPersistance.SetUUID(m_uuid);
     }
 
     if (getOSInfo != nullptr)
@@ -66,15 +65,11 @@ std::vector<std::string> AgentInfo::GetGroups() const
 
 void AgentInfo::SetName(const std::string& name)
 {
-    AgentInfoPersistance agentInfoPersistance;
-    agentInfoPersistance.SetName(name);
     m_name = name;
 }
 
 bool AgentInfo::SetKey(const std::string& key)
 {
-    AgentInfoPersistance agentInfoPersistance;
-
     if (!key.empty())
     {
         if (!ValidateKey(key))
@@ -88,22 +83,16 @@ bool AgentInfo::SetKey(const std::string& key)
         m_key = CreateKey();
     }
 
-    agentInfoPersistance.SetKey(m_key);
-
     return true;
 }
 
 void AgentInfo::SetUUID(const std::string& uuid)
 {
-    AgentInfoPersistance agentInfoPersistance;
-    agentInfoPersistance.SetUUID(uuid);
     m_uuid = uuid;
 }
 
 void AgentInfo::SetGroups(const std::vector<std::string>& groupList)
 {
-    AgentInfoPersistance agentInfoPersistance;
-    agentInfoPersistance.SetGroups(groupList);
     m_groups = groupList;
 }
 
@@ -168,6 +157,15 @@ std::string AgentInfo::GetMetadataInfo(const bool agentIsRegistering) const
     return agentMetadataInfo.dump();
 }
 
+void AgentInfo::Save() const
+{
+    AgentInfoPersistance agentInfoPersistance;
+    agentInfoPersistance.SetName(m_name);
+    agentInfoPersistance.SetKey(m_key);
+    agentInfoPersistance.SetUUID(m_uuid);
+    agentInfoPersistance.SetGroups(m_groups);
+}
+
 std::vector<std::string> AgentInfo::GetActiveIPAddresses(const nlohmann::json& networksJson) const
 {
     std::vector<std::string> ipAddresses;
diff --git a/src/agent/agent_info/tests/agent_info_test.cpp b/src/agent/agent_info/tests/agent_info_test.cpp
@@ -36,6 +36,7 @@ TEST_F(AgentInfoTest, TestPersistedValues)
     agentInfo.SetName("test_name");
     agentInfo.SetKey("4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj");
     agentInfo.SetUUID("test_uuid");
+    agentInfo.Save();
     const AgentInfo agentInfoReloaded;
     EXPECT_EQ(agentInfoReloaded.GetName(), "test_name");
     EXPECT_EQ(agentInfoReloaded.GetKey(), "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj");
@@ -45,25 +46,27 @@ TEST_F(AgentInfoTest, TestPersistedValues)
 TEST_F(AgentInfoTest, TestSetName)
 {
     AgentInfo agentInfo;
+    const std::string oldName = agentInfo.GetName();
     const std::string newName = "new_name";
 
     agentInfo.SetName(newName);
     EXPECT_EQ(agentInfo.GetName(), newName);
 
     const AgentInfo agentInfoReloaded;
-    EXPECT_EQ(agentInfoReloaded.GetName(), newName);
+    EXPECT_EQ(agentInfoReloaded.GetName(), oldName);
 }
 
 TEST_F(AgentInfoTest, TestSetKey)
 {
     AgentInfo agentInfo;
+    const std::string oldKey = agentInfo.GetKey();
     const std::string newKey = "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj";
 
     agentInfo.SetKey(newKey);
     EXPECT_EQ(agentInfo.GetKey(), newKey);
 
     const AgentInfo agentInfoReloaded;
-    EXPECT_EQ(agentInfoReloaded.GetKey(), newKey);
+    EXPECT_EQ(agentInfoReloaded.GetKey(), oldKey);
 }
 
 TEST_F(AgentInfoTest, TestSetBadKey)
@@ -80,12 +83,13 @@ TEST_F(AgentInfoTest, TestSetEmptyKey)
 {
     AgentInfo agentInfo;
     const std::string newKey;
+    const std::string oldKey = agentInfo.GetKey();
 
     agentInfo.SetKey(newKey);
     EXPECT_NE(agentInfo.GetKey(), newKey);
 
     const AgentInfo agentInfoReloaded;
-    EXPECT_NE(agentInfoReloaded.GetKey(), newKey);
+    EXPECT_EQ(agentInfoReloaded.GetKey(), oldKey);
 }
 
 TEST_F(AgentInfoTest, TestSetUUID)
@@ -97,19 +101,20 @@ TEST_F(AgentInfoTest, TestSetUUID)
     EXPECT_EQ(agentInfo.GetUUID(), newUUID);
 
     const AgentInfo agentInfoReloaded;
-    EXPECT_EQ(agentInfoReloaded.GetUUID(), newUUID);
+    EXPECT_NE(agentInfoReloaded.GetUUID(), newUUID);
 }
 
 TEST_F(AgentInfoTest, TestSetGroups)
 {
     AgentInfo agentInfo;
+    const std::vector<std::string> oldGroups = agentInfo.GetGroups();
     const std::vector<std::string> newGroups = {"t_group_1", "t_group_2"};
 
     agentInfo.SetGroups(newGroups);
     EXPECT_EQ(agentInfo.GetGroups(), newGroups);
 
     const AgentInfo agentInfoReloaded;
-    EXPECT_EQ(agentInfoReloaded.GetGroups(), newGroups);
+    EXPECT_EQ(agentInfoReloaded.GetGroups(), oldGroups);
 }
 
 TEST_F(AgentInfoTest, TestLoadMetadataInfoNoSysInfo)
diff --git a/src/agent/src/agent_registration.cpp b/src/agent/src/agent_registration.cpp
@@ -60,12 +60,13 @@ namespace agent_registration
 
         const auto res = httpClient.PerformHttpRequest(reqParams);
 
-        if (res.result() != http::status::ok)
+        if (res.result() != http::status::created)
         {
             std::cout << "Registration error: " << res.result_int() << ".\n";
             return false;
         }
 
+        m_agentInfo.Save();
         return true;
     }
 
diff --git a/src/agent/tests/agent_registration_test.cpp b/src/agent/tests/agent_registration_test.cpp
@@ -71,12 +71,17 @@ TEST_F(RegisterTest, RegistrationTestSuccess)
     AgentInfoPersistance agentInfoPersistance;
     agentInfoPersistance.ResetToDefault();
 
-    registration = std::make_unique<agent_registration::AgentRegistration>(
-        "user", "password", "4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj", "agent_name", std::nullopt);
     SysInfo sysInfo;
     agent = std::make_unique<AgentInfo>([&sysInfo]() mutable { return sysInfo.os(); },
                                         [&sysInfo]() mutable { return sysInfo.networks(); });
 
+    agent->SetKey("4GhT7uFm1zQa9c2Vb7Lk8pYsX0WqZrNj");
+    agent->SetName("agent_name");
+    agent->Save();
+
+    registration = std::make_unique<agent_registration::AgentRegistration>(
+        "user", "password", agent->GetKey(), agent->GetName(), std::nullopt);
+
     MockHttpClient mockHttpClient;
 
     EXPECT_CALL(mockHttpClient, AuthenticateWithUserPassword(testing::_, testing::_, testing::_, testing::_))
@@ -93,7 +98,7 @@ TEST_F(RegisterTest, RegistrationTestSuccess)
                                              bodyJson);
 
     boost::beast::http::response<boost::beast::http::dynamic_body> expectedResponse;
-    expectedResponse.result(boost::beast::http::status::ok);
+    expectedResponse.result(boost::beast::http::status::created);
 
     EXPECT_CALL(mockHttpClient, PerformHttpRequest(testing::Eq(reqParams))).WillOnce(testing::Return(expectedResponse));
 
@@ -145,40 +150,33 @@ TEST_F(RegisterTest, RegistrationFailsIfServerResponseIsNotOk)
     ASSERT_FALSE(res);
 }
 
-TEST_F(RegisterTest, RegistrationTestSuccessWithEmptyKey)
+TEST_F(RegisterTest, RegisteringWithoutAKeyGeneratesOneAutomatically)
 {
     AgentInfoPersistance agentInfoPersistance;
     agentInfoPersistance.ResetToDefault();
 
+    agent = std::make_unique<AgentInfo>();
+    EXPECT_TRUE(agent->GetKey().empty());
+
     registration =
         std::make_unique<agent_registration::AgentRegistration>("user", "password", "", "agent_name", std::nullopt);
-    SysInfo sysInfo;
-    agent = std::make_unique<AgentInfo>([&sysInfo]() mutable { return sysInfo.os(); },
-                                        [&sysInfo]() mutable { return sysInfo.networks(); });
 
     MockHttpClient mockHttpClient;
 
     EXPECT_CALL(mockHttpClient, AuthenticateWithUserPassword(testing::_, testing::_, testing::_, testing::_))
         .WillOnce(testing::Return("token"));
 
-    const auto bodyJson = agent->GetMetadataInfo(true);
-
-    http_client::HttpRequestParams reqParams(boost::beast::http::verb::post,
-                                             "https://localhost:55000",
-                                             "/agents",
-                                             agent->GetHeaderInfo(),
-                                             "token",
-                                             "",
-                                             bodyJson);
-
     boost::beast::http::response<boost::beast::http::dynamic_body> expectedResponse;
-    expectedResponse.result(boost::beast::http::status::ok);
+    expectedResponse.result(boost::beast::http::status::created);
 
-    EXPECT_CALL(mockHttpClient, PerformHttpRequest(testing::Eq(reqParams))).WillOnce(testing::Return(expectedResponse));
+    EXPECT_CALL(mockHttpClient, PerformHttpRequest(testing::_)).WillOnce(testing::Return(expectedResponse));
 
     // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
     const bool res = registration->Register(mockHttpClient);
     ASSERT_TRUE(res);
+
+    agent = std::make_unique<AgentInfo>();
+    EXPECT_FALSE(agent->GetKey().empty());
 }
 
 TEST_F(RegisterTest, RegistrationTestFailWithBadKey)

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,6 @@ AgentInfo::AgentInfo(std::function<nlohmann::json()> getOSInfo, std::function<nl`
`27`	`27`	`if (m_uuid.empty())`
`28`	`28`	`{`
`29`	`29`	`m_uuid = boost::uuids::to_string(boost::uuids::random_generator()());`
`30`		`- agentInfoPersistance.SetUUID(m_uuid);`
`31`	`30`	`}`
`32`	`31`
`33`	`32`	`if (getOSInfo != nullptr)`
`@@ -66,15 +65,11 @@ std::vector<std::string> AgentInfo::GetGroups() const`
`66`	`65`
`67`	`66`	`void AgentInfo::SetName(const std::string& name)`
`68`	`67`	`{`
`69`		`- AgentInfoPersistance agentInfoPersistance;`
`70`		`- agentInfoPersistance.SetName(name);`
`71`	`68`	`m_name = name;`
`72`	`69`	`}`
`73`	`70`
`74`	`71`	`bool AgentInfo::SetKey(const std::string& key)`
`75`	`72`	`{`
`76`		`- AgentInfoPersistance agentInfoPersistance;`
`77`		`-`
`78`	`73`	`if (!key.empty())`
`79`	`74`	`{`
`80`	`75`	`if (!ValidateKey(key))`
`@@ -88,22 +83,16 @@ bool AgentInfo::SetKey(const std::string& key)`
`88`	`83`	`m_key = CreateKey();`
`89`	`84`	`}`
`90`	`85`
`91`		`- agentInfoPersistance.SetKey(m_key);`
`92`		`-`
`93`	`86`	`return true;`
`94`	`87`	`}`
`95`	`88`
`96`	`89`	`void AgentInfo::SetUUID(const std::string& uuid)`
`97`	`90`	`{`
`98`		`- AgentInfoPersistance agentInfoPersistance;`
`99`		`- agentInfoPersistance.SetUUID(uuid);`
`100`	`91`	`m_uuid = uuid;`
`101`	`92`	`}`
`102`	`93`
`103`	`94`	`void AgentInfo::SetGroups(const std::vector<std::string>& groupList)`
`104`	`95`	`{`
`105`		`- AgentInfoPersistance agentInfoPersistance;`
`106`		`- agentInfoPersistance.SetGroups(groupList);`
`107`	`96`	`m_groups = groupList;`
`108`	`97`	`}`
`109`	`98`
`@@ -168,6 +157,15 @@ std::string AgentInfo::GetMetadataInfo(const bool agentIsRegistering) const`
`168`	`157`	`return agentMetadataInfo.dump();`
`169`	`158`	`}`
`170`	`159`
	`160`	`+void AgentInfo::Save() const`
	`161`	`+{`
	`162`	`+ AgentInfoPersistance agentInfoPersistance;`
	`163`	`+ agentInfoPersistance.SetName(m_name);`
	`164`	`+ agentInfoPersistance.SetKey(m_key);`
	`165`	`+ agentInfoPersistance.SetUUID(m_uuid);`
	`166`	`+ agentInfoPersistance.SetGroups(m_groups);`
	`167`	`+}`
	`168`	`+`
`171`	`169`	`std::vector<std::string> AgentInfo::GetActiveIPAddresses(const nlohmann::json& networksJson) const`
`172`	`170`	`{`
`173`	`171`	`std::vector<std::string> ipAddresses;`
Original file line number	Diff line number	Diff line change
`@@ -60,12 +60,13 @@ namespace agent_registration`
`60`	`60`
`61`	`61`	`const auto res = httpClient.PerformHttpRequest(reqParams);`
`62`	`62`
`63`		`- if (res.result() != http::status::ok)`
	`63`	`+ if (res.result() != http::status::created)`
`64`	`64`	`{`
`65`	`65`	`std::cout << "Registration error: " << res.result_int() << ".\n";`
`66`	`66`	`return false;`
`67`	`67`	`}`
`68`	`68`
	`69`	`+ m_agentInfo.Save();`
`69`	`70`	`return true;`
`70`	`71`	`}`
`71`	`72`