From 0f6b4cbc8a8629a3114a91f0e4dddbcc41ad7f93 Mon Sep 17 00:00:00 2001 From: Rob Zimmerman Date: Fri, 7 Jul 2023 17:15:48 +0000 Subject: [PATCH 1/2] adding logging --- chirps/chirps/settings.py | 22 ++++++++++++++++++++++ chirps/scan/tasks.py | 17 ++++++++++------- chirps/target/models.py | 16 +++++++++------- 3 files changed, 41 insertions(+), 14 deletions(-) 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..7dda6886 100644 --- a/chirps/scan/tasks.py +++ b/chirps/scan/tasks.py @@ -1,4 +1,5 @@ """Celery tasks for the scan application.""" +from logging import getLogger import re from celery import shared_task @@ -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..3c16956d 100644 --- a/chirps/target/models.py +++ b/chirps/target/models.py @@ -1,5 +1,5 @@ """Models for the target appliation.""" - +from logging import getLogger import pinecone from django.contrib.auth.models import User from django.db import models @@ -11,7 +11,7 @@ from .custom_fields import CustomEncryptedCharField - +logger = getLogger(__name__) class BaseTarget(PolymorphicModel): """Base class that all targets will inherit from.""" @@ -49,13 +49,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): @@ -101,7 +101,7 @@ def test_connection(self) -> bool: pinecone.deinit() return True except Exception as err: # pylint: disable=broad-exception-caught - print(f"Pinecone connection test failed: {err}") + logger.error('Pinecone connection test failed', extra={'error': err}) return False @@ -119,6 +119,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 +127,7 @@ 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] From 8422ee8ef8fa6c765620711ff91035301a76285b Mon Sep 17 00:00:00 2001 From: Rob Zimmerman Date: Fri, 7 Jul 2023 17:17:50 +0000 Subject: [PATCH 2/2] linting cleanup --- chirps/scan/tasks.py | 8 ++++---- chirps/target/models.py | 11 ++++++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/chirps/scan/tasks.py b/chirps/scan/tasks.py index 7dda6886..cab37417 100644 --- a/chirps/scan/tasks.py +++ b/chirps/scan/tasks.py @@ -1,6 +1,6 @@ """Celery tasks for the scan application.""" -from logging import getLogger import re +from logging import getLogger from celery import shared_task from django.utils import timezone @@ -10,19 +10,19 @@ logger = getLogger(__name__) + @shared_task def scan_task(scan_id): """Main scan task.""" - logger.info('Starting scan', extra={'id': scan_id}) + logger.info('Starting scan', extra={'id': scan_id}) try: scan = Scan.objects.get(pk=scan_id) except Scan.DoesNotExist: 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})'}) + 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 diff --git a/chirps/target/models.py b/chirps/target/models.py index 3c16956d..a888a3c7 100644 --- a/chirps/target/models.py +++ b/chirps/target/models.py @@ -1,5 +1,6 @@ """Models for the target appliation.""" from logging import getLogger + import pinecone from django.contrib.auth.models import User from django.db import models @@ -12,6 +13,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.""" @@ -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,7 +104,7 @@ 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 + except Exception as err: # pylint: disable=broad-exception-caught logger.error('Pinecone connection test failed', extra={'error': err}) return False @@ -130,4 +134,5 @@ def search(self, query: str, max_results: int) -> list[str]: logger.info('Mantium target search complete', extra={'id': self.id}) return documents + targets = [RedisTarget, MantiumTarget, PineconeTarget]