diff --git a/chirps/chirps/settings.py b/chirps/chirps/settings.py index 30d28c88..548043c7 100644 --- a/chirps/chirps/settings.py +++ b/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, + }, + }, +} diff --git a/chirps/scan/tasks.py b/chirps/scan/tasks.py index e1935028..cab37417 100644 --- a/chirps/scan/tasks.py +++ b/chirps/scan/tasks.py @@ -1,5 +1,6 @@ """Celery tasks for the scan application.""" import re +from logging import getLogger from celery import shared_task from django.utils import timezone @@ -7,18 +8,21 @@ 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}) diff --git a/chirps/target/models.py b/chirps/target/models.py index 786457ad..a888a3c7 100644 --- a/chirps/target/models.py +++ b/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,7 +83,7 @@ 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]: @@ -87,7 +91,7 @@ def search(self, query: str, max_results: int) -> list[str]: 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,6 +123,7 @@ 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) @@ -126,6 +131,8 @@ def search(self, query: str, max_results: int) -> list[str]: 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]