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()