Merge pull request #49 from mantiumai/zim/37

Adding logging configuration

chirps/chirps/settings.py

@@ -151,3 +151,25 @@
     raise Exception('FERNET_KEY environment variable is not set') # pylint: disable=broad-exception-raised
 
 FERNET_KEY = os.getenv('FERNET_KEY')
+
+# LOGGING Configuration Options
+LOGGING = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "handlers": {
+        "console": {
+            "class": "logging.StreamHandler",
+        },
+    },
+    "root": {
+        "handlers": ["console"],
+        "level": "WARNING",
+    },
+    "loggers": {
+        "django": {
+            "handlers": ["console"],
+            "level": os.getenv("DJANGO_LOG_LEVEL", "INFO"),
+            "propagate": False,
+        },
+    },
+}

chirps/scan/tasks.py

@@ -1,24 +1,28 @@
 """Celery tasks for the scan application."""
 import re
+from logging import getLogger
 
 from celery import shared_task
 from django.utils import timezone
 from target.models import BaseTarget
 
 from .models import Finding, Result, Scan
 
+logger = getLogger(__name__)
+
 
 @shared_task
 def scan_task(scan_id):
     """Main scan task."""
 
-    print(f'Running a scan {scan_id}')
+    logger.info('Starting scan', extra={'id': scan_id})
+
     try:
         scan = Scan.objects.get(pk=scan_id)
     except Scan.DoesNotExist:
-        error = f'Scan {scan_id} does not exist'
-        print(error)
-        scan_task.update_state(state='FAILURE', meta={'error': error})
+        logger.error('Scan record not found', extra={'id': scan_id})
+
+        scan_task.update_state(state='FAILURE', meta={'error': f'Scan record not found ({scan_id})'})
         return
 
     # Need to perform a secondary query in order to fetch the derrived class
@@ -27,8 +31,7 @@ def scan_task(scan_id):
 
     # Now that we have the derrived class, call its implementation of search()
     for rule in scan.plan.rules.all():
-        print(f'Running rule {rule}')
-
+        logger.info('Starting rule evaluation', extra={'id': rule.id})
         results = target.search(query=rule.query_string, max_results=100)
 
         for text in results:
@@ -47,4 +50,4 @@ def scan_task(scan_id):
     # Persist the completion time of the scan
     scan.finished_at = timezone.now()
     scan.save()
-    print('Saved scan results')
+    logger.info('Scan complete', extra={'id': scan_id})

chirps/target/models.py

@@ -1,4 +1,5 @@
 """Models for the target appliation."""
+from logging import getLogger
 
 import pinecone
 from django.contrib.auth.models import User
@@ -11,6 +12,8 @@
 
 from .custom_fields import CustomEncryptedCharField
 
+logger = getLogger(__name__)
+
 
 class BaseTarget(PolymorphicModel):
     """Base class that all targets will inherit from."""
@@ -33,6 +36,7 @@ def __str__(self) -> str:
         """String representation of this model."""
         return str(self.name)
 
+
 class RedisTarget(BaseTarget):
     """Implementation of a Redis target."""
 
@@ -49,13 +53,13 @@ class RedisTarget(BaseTarget):
 
     def search(self, query: str, max_results: int) -> str:
         """Search the Redis target with the specified query."""
-        print('Starting RedisTarget search')
-        print('Converting search query into an embedding vector')
-        print('RedisTarget search copmlete')
+        logger.error('RedisTarget search not implemented')
+        raise NotImplementedError
 
     def test_connection(self) -> bool:
         """Ensure that the Redis target can be connected to."""
-        return True
+        logger.error('RedisTarget search not implemented')
+        raise NotImplementedError
 
 
 class PineconeTarget(BaseTarget):
@@ -79,15 +83,15 @@ def decrypted_api_key(self):
                 decrypted_value = self.api_key
                 return decrypted_value
             except UnicodeDecodeError:
-                return "Error: Decryption failed"
+                return 'Error: Decryption failed'
         return None
 
     def search(self, query: str, max_results: int) -> list[str]:
         """Search the Pinecone target with the specified query."""
         pinecone.init(api_key=self.api_key, environment=self.environment)
 
         # Assuming the query is converted to a vector of the same dimension as the index. We should re-visit this.
-        query_vector = convert_query_to_vector(query) # pylint: disable=undefined-variable
+        query_vector = convert_query_to_vector(query)   # pylint: disable=undefined-variable
 
         # Perform search on the Pinecone index
         search_results = pinecone.fetch(index_name=self.index_name, query_vector=query_vector, top_k=max_results)
@@ -100,8 +104,8 @@ def test_connection(self) -> bool:
             pinecone.init(api_key=self.api_key, environment=self.environment)
             pinecone.deinit()
             return True
-        except Exception as err: # pylint: disable=broad-exception-caught
-            print(f"Pinecone connection test failed: {err}")
+        except Exception as err:   # pylint: disable=broad-exception-caught
+            logger.error('Pinecone connection test failed', extra={'error': err})
             return False
 
 
@@ -119,13 +123,16 @@ class MantiumTarget(BaseTarget):
     html_description = 'Mantium Knowledge Vault'
 
     def search(self, query: str, max_results: int) -> list[str]:
+        logger.info('Starting Mantium Target search', extra={'id': self.id})
         client = MantiumClient(client_id=self.client_id, client_secret=self.client_secret)
         apps_api = ApplicationsApi(client)
 
         query_request = {'query': query}
         results = apps_api.query_application(self.app_id, query_request)
 
         documents = [doc['content'] for doc in results['documents']]
+        logger.info('Mantium target search complete', extra={'id': self.id})
         return documents
 
+
 targets = [RedisTarget, MantiumTarget, PineconeTarget]