Merge branch 'updated-error-handling' into 'master'

Updated error handling and tests

See merge request medina/evidence-collector!10

.gitlab-ci.yml

@@ -24,11 +24,14 @@ test:
   stage: test
   script:
     - apk add bash
-    - docker stop $SERVICE || true && docker rm $SERVICE || true
-    - docker run --env-file .env --name $SERVICE -d $REGISTRY/medina/$SERVICE:$VERSION
-    - sleep 5
+    - docker network create test-ec
+    - docker run --rm --network=test-ec --env-file .env --name $SERVICE -d $REGISTRY/medina/$SERVICE:$VERSION
+    - docker run --rm --network=test-ec toschneck/wait-for-it $SERVICE:7890 -t 240
     - bash test/test.sh
-    - docker stop $SERVICE && docker container rm $SERVICE
+  after_script:
+    - SERVICE=$(grep SERVICE MANIFEST | cut -d '=' -f2)
+    - docker kill $SERVICE || docker network rm test-ec
+    - docker network rm test-ec
 
 push:
   stage: push

Dockerfile

@@ -9,6 +9,6 @@ RUN pip3 install -r requirements.txt
 
 COPY . .
 
-RUN apt-get update && apt-get install -y redis-server
+RUN apt-get update && apt-get install -y redis-server netcat
 
-ENTRYPOINT ["./entrypoint.sh"]
\ No newline at end of file
+ENTRYPOINT ["./entrypoint.sh"]

MANIFEST

-VERSION=v0.0.11
+VERSION=v0.0.12
 SERVICE=evidence-collector

Makefile

@@ -2,4 +2,4 @@ build:
 	docker build -t evidence-collector .
 
 run:
-	docker run --env-file .env  evidence-collector
\ No newline at end of file
+	docker run --env-file .env evidence-collector

entrypoint.sh

@@ -8,4 +8,8 @@ rqscheduler --host $redis_host --port $redis_port &
 
 python3 -m scheduler.scheduler
 
-tail -f /var/log/evidence_collector.log
\ No newline at end of file
+# open a listener on port 7890 for 1 second
+# only for testing - CI script contains a wait-for-it that binds to this port
+nc -l -p 7890 -w 1
+
+tail -f /var/log/evidence_collector.log

forward_evidence/clouditor_authentication.py

@@ -30,6 +30,8 @@ class ClouditorAuthentication(object):
                 urllib3.exceptions.MaxRetryError, requests.exceptions.ConnectionError) as err:
             self.logger.error(err)
             self.logger.error("Clouditor OAuth2 token endpoint not available")
+            self.__access_token = None
+            self.__token_expiration_time = None
         else:
             token = json.loads(access_token_response.text)
 
@@ -41,7 +43,7 @@ class ClouditorAuthentication(object):
 
     def get_token(self):
         # In practice this condition isn't even needed as every scheduled job creates new ClouditorAuthentication object and acquires new token.
-        if (datetime.utcnow() > self.__token_expiration_time):
+        if (self.__token_expiration_time != None and datetime.utcnow() > self.__token_expiration_time):
             self.logger.debug("OAuth2 token expired.")
             self.request_token()
 

forward_evidence/forward_evidence.py

@@ -14,7 +14,10 @@ class ForwardEvidence(object):
 
     def send_evidence(self, assessevidencerequest, token):
         try:
-            metadata = [('authorization', 'Bearer ' + token)]
+            if token is not None:
+                metadata = [('authorization', 'Bearer ' + token)]
+            else:
+                metadata = None
 
             response = self.stub.AssessEvidence(assessevidencerequest, metadata=metadata)
             self.logger.info('gRPC evidence forwarded: ' + str(response))

test/test.sh

@@ -6,6 +6,7 @@ redis2="Ready to accept connections"
 scheduler="Registering birth"
 worker1="Worker rq:worker:"
 worker2="Listening on "
+oauth2token="Clouditor OAuth2 token endpoint not available"
 
 if ! [[ $logs =~ $redis1 ]]
     then
@@ -36,3 +37,9 @@ if ! [[ $logs =~ $worker2 ]]
         echo "Redis worker not started" 1>&2
         exit 1
 fi
+
+if ! [[ $logs =~ $oauth2token ]]
+    then
+        echo "OAuth2 token authentication not working" 1>&2
+        exit 1
+fi

wazuh_evidence_collector/checker.py

 from wazuh_evidence_collector.wazuh_client import WazuhClient
-from elasticsearch import Elasticsearch
+import elasticsearch
+import urllib3
 from elasticsearch_dsl import Search
 
 class Checker:
-    def __init__(self, wc, es):
+    def __init__(self, wc, es, logger):
         self.wc = wc
         self.es = es
+        self.logger = logger
         
     # Check if syscheck enabled
     def check_syscheck(self, agent):
@@ -92,8 +94,15 @@ class Checker:
             .query("match", rule__description="Clamd restarted") \
             .query("match", agent__id=agent[0])
         
-        body = s.execute().to_dict()
-        
+        try:
+            body = s.execute().to_dict()
+        except (elasticsearch.exceptions.ConnectionError, TimeoutError, urllib3.exceptions.NewConnectionError, 
+        		urllib3.exceptions.MaxRetryError) as err:
+            self.logger.error(err)
+            self.logger.error("Elasticsearch not available")
+
+            return None, False
+
         measurement_result = len(body['hits']['hits']) > 0
         
         return body, measurement_result

wazuh_evidence_collector/wazuh_client.py

@@ -3,12 +3,13 @@ import urllib3
 
 class WazuhClient:
 	
-	def __init__(self, ip, port, username, password):
+	def __init__(self, ip, port, username, password, logger):
 		self._ip = ip
 		self._port = port
 		self._username = username
 		self._password = password
 		self._auth_token = None
+		self.logger = logger
 	
 	def req(self, method, resource, data=None, headers={}, auth_retry=True):
 		# TODO: add cert verification
@@ -19,7 +20,13 @@ class WazuhClient:
 		if self._auth_token:
 			headers['Authorization'] = 'Bearer %s' % self._auth_token
 		
-		resp = c.request(method, url, headers=headers, body=data)
+		try:
+			resp = c.request(method, url, headers=headers, body=data)
+		except (TimeoutError, urllib3.exceptions.NewConnectionError, 
+        		urllib3.exceptions.MaxRetryError) as err:
+			self.logger.error(err)
+			self.logger.error("Wazuh manager not available")
+
 		if resp.status == 401:
 			if not auth_retry:
 				raise Exception("Authentication Error")

wazuh_evidence_collector/wazuh_evidence_collector.py

@@ -27,7 +27,7 @@ ELASTIC_USERNAME = os.environ.get("elastic_username")
 ELASTIC_PASSWORD = os.environ.get("elastic_password")
 
 if not DEMO:
-    wc = WazuhClient(WAZUH_HOST, WAZUH_PORT, WAZUH_USERNAME, WAZUH_PASSWORD)
+    wc = WazuhClient(WAZUH_HOST, WAZUH_PORT, WAZUH_USERNAME, WAZUH_PASSWORD, LOGGER)
 
     es = Elasticsearch(
             ELASTIC_HOST,
@@ -69,7 +69,7 @@ def main():
 
 # Wrapper function that runs all the checks (for every manager/agent)
 def run_collector():
-    checker = Checker(wc, es) if not DEMO else DemoChecker()
+    checker = Checker(wc, es, LOGGER) if not DEMO else DemoChecker()
 
     # Get list of all agent ids (including manager's)
     def get_agents(wc):