add port and IP for API

deps/wazuh_testing/wazuh_testing/tools/performance/statistic.py

@@ -1,6 +1,3 @@
-
-
-
 # Copyright (C) 2015-2021, Wazuh Inc.
 # Created by Wazuh, Inc. <info@wazuh.com>.
 # This program is free software; you can redistribute it and/or modify it under the terms of GPLv2
@@ -21,10 +18,6 @@
 
 urllib3.disable_warnings()
 
-API_URL="https://localhost:55000"
-DAEMONS_ENDPOINT="/manager/daemons/stats?daemons_list=wazuh-analysisd,wazuh-remoted,wazuh-db"
-TOKEN_ENDPOINT="/security/user/authenticate"
-
 logger = logging.getLogger('wazuh-statistics-monitor')
 
 
@@ -44,6 +37,8 @@ class StatisticMonitor:
         target (str, optional): target file to monitor.
         dst_dir (str, optional): path to store the file.
         use_api (bool, optional): Determine if the API should be used to collect the data. Default False.
+        ip (str, optional): IP address of the API. Default localhost.
+        port (str, optional): Port of the API. Default 55000.
 
     Attributes:
         event (thread.Event): thread Event used to control the scans.
@@ -55,16 +50,19 @@ class StatisticMonitor:
         parse_json (bool): Determine if the file is a JSON file. Default False.
     """
 
-    def __init__(self, target='agent', time_step=5, dst_dir=gettempdir(), use_api=False):
+    def __init__(self, target='agent', time_step=5, dst_dir=gettempdir(), use_api=False, ip = 'localhost', port = '55000'):
         self.event = None
         self.thread = None
         self.time_step = time_step
         self.target = target
         self.dst_dir = dst_dir
         self.use_api = use_api
+        self.ip = ip
+        self.port = port
+        self.daemon = None
         self.parse_json = False
 
-        if self.use_api == True and self.target != 'analysis_events' and self.target != 'remote' and self.target != 'wazuhdb':
+        if self.use_api == True and self.target != 'analysis' and self.target != 'remote' and self.target != 'wazuhdb':
             self.use_api = False
 
         if self.target == 'agent':
@@ -74,10 +72,13 @@ def __init__(self, target='agent', time_step=5, dst_dir=gettempdir(), use_api=Fa
             self.parse_json = True
         elif self.target == 'remote':
             self.statistics_file = tls.REMOTE_STATISTICS_FILE
-        elif self.target == 'analysis_state':
+            self.daemon = 'wazuh-remoted'
+        elif self.target == 'analysis':
             self.statistics_file = tls.ANALYSIS_STATISTICS_FILE
-        elif self.target == 'analysis_events' or self.target == 'wazuhdb':
+            self.daemon = 'wazuh-analysisd'
+        elif self.target == 'wazuhdb':
             self.use_api = True
+            self.daemon = 'wazuh-db'
         else:
             raise ValueError(f'The target {self.target} is not a valid one.')
 
@@ -140,22 +141,26 @@ def _parse_state_file(self):
             logger.error(f'Exception with {self.statistics_file} | {str(e)}')
 
 
-    def _parse_api_data(self):
+    def _parse_api_data(self, ip, port):
         """Read the data from the statistics file generated by Wazuh API."""
 
+        API_URL = f"https://{ip}:{port}"
+        DAEMONS_ENDPOINT= f"/manager/daemons/stats?daemons_list={self.daemon}"
+        TOKEN_ENDPOINT="/security/user/authenticate"
+
         logging.info("Getting statistics data from API for {}".format(self.target))
 
         response = requests.get(API_URL + TOKEN_ENDPOINT, verify=False, auth=requests.auth.HTTPBasicAuth("wazuh", "wazuh"))
         if response.status_code != 200:
-            logging.info("Retrying get API data, status code {}".format(response.status_code))
-            return self._parse_api_data()
+            logging.error("Retrying get API data, status code {}".format(response.status_code))
+            return self._parse_api_data(ip, port)
 
         daemons_response = requests.get(API_URL + DAEMONS_ENDPOINT, verify=False, headers={'Authorization': 'Bearer ' + response.json()['data']['token']})
         if daemons_response.status_code != 200:
-            logging.info("Retrying get API data, status code {}".format(response.status_code))
-            return self._parse_api_data()
+            logging.error("Retrying get API data, status code {}".format(response.status_code))
+            return self._parse_api_data(ip, port)
 
-        data = daemons_response.json()['data']['affected_items']
+        data = daemons_response.json()['data']['affected_items'][0]
         self._write_csv(data, self.target, self.csv_file)
 
 
@@ -167,10 +172,8 @@ def _write_csv(self, data, target, csv_file):
             target (string): specifies which CSV must be generated
             csv_file (string): path to the CSV file.
         """
-        if target == "analysis_state":
-            csv_header = headers.analysisd_header
-        elif target == "analysis_events":
-            csv_header = headers.analysisd_events_header
+        if target == "analysis":
+            csv_header = headers.analysisd_header if self.use_api == False else headers.analysisd_events_header
         elif target == "logcollector":
             csv_header = headers.logcollector_header
         elif target == "remote":
@@ -190,21 +193,12 @@ def _write_csv(self, data, target, csv_file):
             timestamp = datetime.fromtimestamp(time()).strftime('%Y-%m-%d %H:%M:%S')
 
             if self.use_api:
-
-                ## Get data from API response
-                if target == "analysis_events":
-                    data = data[0]
-                elif target == "remote":
-                    data = data[1]
-                elif target == "wazuhdb":
-                    data = data[2]
-
                 format = r"%Y-%m-%dT%H:%M:%S+%f:00"
                 datetime_timestamp = datetime.strptime(data['timestamp'], format)
                 datetime_uptime = datetime.strptime(data['uptime'], format)
                 interval =  (datetime_timestamp - datetime_uptime).total_seconds()
 
-                if target == "analysis_events":                   
+                if target == "analysis":                   
                     metrics = data['metrics']
                     decoded = metrics['events']['received_breakdown']['decoded_breakdown']
                     decoded_modules = decoded['modules_breakdown']
@@ -488,7 +482,7 @@ def _write_csv(self, data, target, csv_file):
                         execution_breakdown['wazuhdb']                                                              # 148
                     ))
             else:
-                if target == "analysis_state":
+                if target == "analysis":
                     logger.info("Writing analysisd.state info to {}.".format(csv_file))
                     log.write(("{0},{1},{2},{3},{4},{5},{6},{7},{8},{9},{10},{11},{12},{13},{14},{15},{16},{17},{18},{19},{20},"+
                         "{21},{22},{23},{24},{25},{26},{27},{28},{29},{30},{31},{32},{33},{34},{35},{36},{37},{38},{39},{40},{41},{42}\n")
@@ -565,7 +559,7 @@ def _monitor_stats(self):
         """Read the .state files and log the data into a CSV file."""
         while not self.event.is_set():
             if self.use_api:
-                self._parse_api_data()
+                self._parse_api_data(self.ip, self.port)
             else:
                 self._parse_state_file()
             sleep(self.time_step)