Merge pull request #2295 from wazuh/feature/10932-scripts-simulate-cl…

…uster-load

Add scripts to add agents to client.keys, create agent-groups and unsynchronize agents

deps/wazuh_testing/setup.py

@@ -36,7 +36,10 @@
     'simulate-api-load=wazuh_testing.scripts.simulate_api_load:main',
     'wazuh-log-metrics=wazuh_testing.scripts.wazuh_log_metrics:main',
     'qa-docs=wazuh_testing.scripts.qa_docs:main',
-    'qa-ctl=wazuh_testing.scripts.qa_ctl:main'
+    'qa-ctl=wazuh_testing.scripts.qa_ctl:main',
+    'add-agents-client-keys=wazuh_testing.scripts.add_agents_client_keys:main',
+    'add-agents-to-default-group=wazuh_testing.scripts.add_agents_to_default_group:main',
+    'unsync-agents=wazuh_testing.scripts.unsync_agents:main'
 ]
 
 

deps/wazuh_testing/wazuh_testing/scripts/add_agents_client_keys.py

@@ -0,0 +1,31 @@
+import sys
+
+from wazuh_testing.tools import CLIENT_KEYS_PATH
+
+
+def main():
+    """Add fake agents to client.keys. To use the script, pass two arguments indicating the first agent ID and the last
+    agent ID from the range of agents to be added.
+
+    The agents added will have ID={agent_id}, name=new_agent_{agent_id}, address=any; and password={agent_id}.
+
+    This script must be used in the Wazuh master node.
+    """
+    if len(sys.argv) != 3:
+        print(f"add_agents_client_keys.py <first_id> <last_id> (you used {' '.join(sys.argv)})")
+        exit(1)
+
+    first_id = min(int(sys.argv[1]), int(sys.argv[2]))
+    last_id = max(int(sys.argv[1]), int(sys.argv[2]))
+
+    agents_list = [str(agent_id).zfill(3) for agent_id in range(first_id, last_id + 1)]
+
+    with open(file=CLIENT_KEYS_PATH, mode='a') as f:
+        for agent_id in agents_list:
+            f.write(f"{agent_id} new_agent_{agent_id} any {agent_id}\n")
+            f.flush()  # Avoid bytes staying in the buffer until the loop has finished
+    exit(0)
+
+
+if __name__ == "__main__":
+    main()

deps/wazuh_testing/wazuh_testing/scripts/add_agents_to_default_group.py

@@ -0,0 +1,39 @@
+import os
+import sys
+
+from wazuh_testing.tools import CLIENT_KEYS_PATH
+
+
+def main():
+    """Add agents to default group. To use the script, pass two arguments indicating the first agent ID and the last
+    agent ID from the range of agents to be added to the default group.
+
+    The agents for which agent-groups will be created will be the intersection of the agent list generated and the list
+    of agents whose ID is in the client.keys file.
+
+    This script must be used in a Wazuh worker node.
+    """
+    if len(sys.argv) != 3:
+        print(f"add_agents_to_default_group.py <first_id> <last_id> (you used {' '.join(sys.argv)})")
+        exit(1)
+
+    first_id = min(int(sys.argv[1]), int(sys.argv[2]))
+    last_id = max(int(sys.argv[1]), int(sys.argv[2]))
+
+    agents_list = [str(agent_id).zfill(3) for agent_id in range(first_id, last_id + 1)]
+
+    with open(file=CLIENT_KEYS_PATH, mode='r') as f:
+        agents_in_client_keys = f.read().split('\n')[:-1]
+    available_agents = [agent.split()[0] for agent in agents_in_client_keys]
+
+    for agent_id in set(agents_list).intersection(available_agents):
+        agent_group_file = f"/var/ossec/queue/agent-groups/{agent_id}"
+        if not os.path.exists(agent_group_file):
+            with open(file=agent_group_file, mode='w') as f:
+                f.write('default')
+
+    exit(0)
+
+
+if __name__ == '__main__':
+    main()

deps/wazuh_testing/wazuh_testing/scripts/unsync_agents.py

@@ -0,0 +1,79 @@
+import socket
+import struct
+import sys
+import time
+
+
+def main():
+    """Desynchronize agents in the worker's global.db. To use the script, pass three arguments indicating the first
+    agent ID, the last agent ID from the range of agents to be added to the default group; and the node name.
+
+    The script updates the global.db agent table entries where ID is one in the range specified. This update includes
+    setting the node_name and the agent version. After that, each agent is marked as required to be synchronized
+    (synreq) every 10 seconds.
+
+    This script must be used in a Wazuh worker node.
+    """
+
+    def send_msg(msg):
+        """Send message to a socket.
+
+        Args:
+            msg (str): Message to be sent to the socket.
+        """
+        msg = struct.pack('<I', len(msg)) + msg.encode()
+
+        # Send msg
+        sock.send(msg)
+
+        # Receive response
+        data = sock.recv(4)
+        data_size = struct.unpack('<I', data[0:4])[0]
+        data = sock.recv(data_size).decode(encoding='utf-8', errors='ignore').split(" ", 1)
+
+        return data
+
+    if len(sys.argv) != 4:
+        msg = f"unsync.py <first_id> <last_id> <node_name> (you used {' '.join(sys.argv)})"
+        print(msg)
+        exit(1)
+
+    first_id = min(int(sys.argv[1]), int(sys.argv[2]))
+    last_id = max(int(sys.argv[1]), int(sys.argv[2]))
+    node_name = sys.argv[3]
+
+    ADDR = '/var/ossec/queue/db/wdb'
+
+    while True:
+        try:
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            sock.connect(ADDR)
+            msg = f'global sql UPDATE agent SET node_name = "{node_name}", version="Wazuh v4.0.0" ' \
+                  f'where id>{first_id} and id<={last_id}'
+            print(f"Updating node_name ({node_name}) and version of the agents: {send_msg(msg)}")
+            sock.close()
+            break
+        except Exception as e:
+            print(f"Could not find wdb socket: {e}. Retrying in 10 seconds...")
+            time.sleep(10)
+
+    while True:
+        try:
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            sock.connect(ADDR)
+            msg = f'global sql UPDATE agent SET sync_status="syncreq", last_keepalive="{int(time.time())}", ' \
+                  f'connection_status="active" where id>{first_id} and id<={last_id}'
+            print(f"Updating sync_status of agents between {first_id} and {last_id}: {send_msg(msg)}")
+            sock.close()
+            time.sleep(10)
+        except KeyboardInterrupt:
+            print("Closing socket")
+            sock.close()
+            exit(1)
+        except Exception as e:
+            print(f"An exception was raised: {e}. Retrying in 10 seconds...")
+            time.sleep(10)
+
+
+if __name__ == '__main__':
+    main()