diff --git a/stixify/web/serializers.py b/stixify/web/serializers.py index 1c0019e..ca7f72c 100644 --- a/stixify/web/serializers.py +++ b/stixify/web/serializers.py @@ -1,5 +1,7 @@ from enum import StrEnum, auto import logging +import re +from urllib.parse import urlparse from rest_framework import serializers, validators from dogesec_commons.utils.serializers import JSONSchemaSerializer @@ -110,7 +112,8 @@ class FileSerializer(serializers.ModelSerializer): ai_incident_classification = serializers.ListField(required=False, read_only=True, allow_null=True) summary = serializers.CharField(read_only=True, required=False, allow_null=True) archived_pdf = serializers.FileField(use_url=True, read_only=True, allow_null=True) - sources = CharacterSeparatedField(required=False, allow_null=True, help_text="You can use this to add one or more sources to the `external_references` property of the Report object created. Useful for tracking locations (i.e. URLs) where the report was sourced.", child=serializers.CharField(max_length=1024)) + labels = CharacterSeparatedField(required=False, allow_null=True, help_text="Labels must contain only lowercase letters, numbers, and hyphens. Separate multiple labels with commas.", child=serializers.SlugField(max_length=256), max_length=32) + sources = CharacterSeparatedField(required=False, allow_null=True, help_text="You can use this to add one or more sources to the `external_references` property of the Report object created. Sources must be valid URLs. Separate multiple sources with commas.", child=serializers.URLField(max_length=1024), max_length=32) confidence = serializers.IntegerField( required=False, allow_null=True, @@ -153,6 +156,7 @@ class FilePatchSerializer(FileSerializer): class Meta: model = File fields = ["name", "labels", "sources"] + def validate(self, attrs): if not attrs: raise ValidationError( diff --git a/stixify/worker/tasks.py b/stixify/worker/tasks.py index 3f314ba..fdff9a4 100644 --- a/stixify/worker/tasks.py +++ b/stixify/worker/tasks.py @@ -3,8 +3,10 @@ import os from pathlib import Path import profile +import time import uuid from django.utils import timezone +from txt2stix import txt2stixBundler from stixify.web.models import Job, File from stixify.web import models from celery import shared_task @@ -16,6 +18,7 @@ from django.core.files.base import File as DjangoFile from django.core.files.base import File as DjangoFile from django.db import transaction +from django.core.cache import cache import stix2 from stixify.worker import helpers, pdf_converter @@ -24,6 +27,42 @@ POLL_INTERVAL = 1 +ARANGO_UPLOAD_COUNTER_KEY = "arango_upload_active_count" +MAX_CONCURRENT_UPLOADS = 1 +LOCK_TIMEOUT = 300 + + +def acquire_upload_lock(job_id, wait_timeout=LOCK_TIMEOUT): + lock_key = f"arango_upload_lock:{job_id}" + logging.info(f"Attempting to acquire upload lock: {lock_key}") + lock_acquired_at = cache.get(lock_key) + + if lock_acquired_at is not None: + raise RuntimeError(f"Upload lock already held for job {job_id}") + + start_time = time.time() + while True: + active_count = cache.get(ARANGO_UPLOAD_COUNTER_KEY, 0) + if active_count < MAX_CONCURRENT_UPLOADS: + cache.set(ARANGO_UPLOAD_COUNTER_KEY, active_count + 1, LOCK_TIMEOUT) + cache.set(lock_key, time.time(), LOCK_TIMEOUT) + logging.info(f"Acquired upload lock for job {job_id} (active: {active_count + 1}/{MAX_CONCURRENT_UPLOADS})") + return + + if time.time() - start_time > wait_timeout: + raise TimeoutError(f"Timeout waiting for arango upload slot after {wait_timeout}s") + + time.sleep(0.1) + + +def release_upload_lock(job_id): + lock_key = f"arango_upload_lock:{job_id}" + if cache.get(lock_key) is not None: + cache.delete(lock_key) + active_count = cache.get(ARANGO_UPLOAD_COUNTER_KEY, 0) + if active_count > 0: + cache.set(ARANGO_UPLOAD_COUNTER_KEY, active_count - 1, LOCK_TIMEOUT) + logging.info(f"Released upload lock for job {job_id} (active: {max(0, active_count - 1)}/{MAX_CONCURRENT_UPLOADS})") def new_task(job: Job): @@ -43,6 +82,25 @@ def create_reprocessing_job(file: File, options: dict = None): new_task(job) return job +def _process_file(processor, job, file): + skip_extraction = bool((job.extra or {}).get("skip_extraction")) + is_reprocess = job.type == models.JobType.REPROCESS_POSTS + + if is_reprocess and skip_extraction: + processor.output_md = file.markdown_file.open().read().decode() + if not file.txt2stix_data: + raise Exception("no existing extraction data to use for reprocess with skip_extraction=true") + txt2stix_data = Txt2StixData.model_validate(file.txt2stix_data) + processor.txt2stix(txt2stix_data) + else: + logging.info(f"running file2txt on {processor.task_name}") + processor.file2txt() + logging.info(f"running txt2stix on {processor.task_name}") + processor.txt2stix() + + processor.write_bundle(processor.bundler) + + @shared_task def process_post(job_id, *args): job = Job.objects.get(id=job_id) @@ -88,23 +146,18 @@ def process_post(job_id, *args): processor.setup( report_prop=report_props, extra=dict(_stixify_file_id=str(file.id)) ) - skip_extraction = bool((job.extra or {}).get("skip_extraction")) - # remove existing values for this file that are not in the new upload (handles deletions and modifications) models.ObjectValue.objects.filter(file_id=file.id).delete() - if job.type == models.JobType.REPROCESS_POSTS and skip_extraction: - processor.output_md = file.markdown_file.open().read().decode() - txt2stix_data = None - if not file.txt2stix_data: - raise Exception("no existing extraction data to use for reprocess with skip_extraction=true") - txt2stix_data = Txt2StixData.model_validate(file.txt2stix_data) - processor.txt2stix(txt2stix_data) - processor.write_bundle(processor.bundler) + _process_file(processor, job, file) + + acquire_upload_lock(job.id) + try: + logging.info(f"uploading {processor.task_name} to arangodb via stix2arango") processor.upload_to_arango() - else: - processor.process() + finally: + release_upload_lock(job.id) - with transaction.atomic(): # revert to old file if something goes wrong during processing + with transaction.atomic(): new_profile_id = (job.extra or {}).get("profile_id") if new_profile_id: file.profile_id = new_profile_id @@ -112,9 +165,9 @@ def process_post(job_id, *args): file.set_txt2stix_data(processor.txt2stix_data) file.create_embedding(include_non_incident=settings.CREATE_EMBEDDING_INCLUDE_NON_INCIDENT) - if job.type == models.JobType.IMPORT_FILE: # only update files for import jobs, reprocess jobs should keep the same file references + if job.type == models.JobType.IMPORT_FILE: file.markdown_file.save("markdown.md", processor.md_file.open(), save=True) - models.FileImage.objects.filter(report=file).delete() # remove old references + models.FileImage.objects.filter(report=file).delete() for image in processor.md_images: models.FileImage.objects.create( diff --git a/tests/src/test_tasks.py b/tests/src/test_tasks.py index 88cc078..bb5f6d7 100644 --- a/tests/src/test_tasks.py +++ b/tests/src/test_tasks.py @@ -183,7 +183,7 @@ def test_process_post_reprocess_skip_extraction_uses_existing_data( ): mock_stixify_processor_cls.return_value = fake_stixifier_processor new_task(stixify_reprocess_job) - fake_stixifier_processor.process.assert_not_called() + fake_stixifier_processor.file2txt.assert_not_called() fake_stixifier_processor.txt2stix.assert_called_once() fake_stixifier_processor.write_bundle.assert_called_once() fake_stixifier_processor.upload_to_arango.assert_called_once() @@ -191,6 +191,62 @@ def test_process_post_reprocess_skip_extraction_uses_existing_data( mock_create_embedding.assert_called_once() + +@pytest.mark.django_db +def test_process_post_reprocess_skip_extraction_acquires_lock( + stixify_reprocess_job, fake_stixifier_processor +): + from django.core.cache import cache + from stixify.worker.tasks import ARANGO_UPLOAD_COUNTER_KEY + + file = stixify_reprocess_job.file + file.markdown_file.save("test.md", io.BytesIO(b"test content")) + file.save(update_fields=["markdown_file", "txt2stix_data"]) + stixify_reprocess_job.extra = {"skip_extraction": True} + stixify_reprocess_job.save(update_fields=["extra"]) + + cache.clear() + + with ( + patch("stixify.worker.tasks.StixifyProcessor") as mock_stixify_processor_cls, + patch.object(models.File, "create_embedding") as mock_create_embedding, + ): + mock_stixify_processor_cls.return_value = fake_stixifier_processor + process_post.si(stixify_reprocess_job.id).delay() + + lock_key = f"arango_upload_lock:{stixify_reprocess_job.id}" + assert cache.get(lock_key) is None, "Lock should be released after upload" + assert cache.get(ARANGO_UPLOAD_COUNTER_KEY, 0) == 0, "Counter should be 0 after upload" + fake_stixifier_processor.upload_to_arango.assert_called_once() + + +@pytest.mark.django_db +def test_process_post_concurrent_uploads_limited( + stixify_reprocess_job, fake_stixifier_processor +): + from django.core.cache import cache + from stixify.worker.tasks import ARANGO_UPLOAD_COUNTER_KEY, MAX_CONCURRENT_UPLOADS + + file = stixify_reprocess_job.file + file.markdown_file.save("test.md", io.BytesIO(b"test content")) + file.save(update_fields=["markdown_file", "txt2stix_data"]) + + cache.clear() + + with ( + patch("stixify.worker.tasks.StixifyProcessor") as mock_stixify_processor_cls, + patch.object(models.File, "create_embedding") as mock_create_embedding, + ): + mock_stixify_processor_cls.return_value = fake_stixifier_processor + + cache.set(ARANGO_UPLOAD_COUNTER_KEY, MAX_CONCURRENT_UPLOADS, 300) + + from stixify.worker.tasks import acquire_upload_lock + with pytest.raises(TimeoutError): + acquire_upload_lock(stixify_reprocess_job.id, wait_timeout=0.1) + + + @pytest.mark.django_db def test_process_post_reprocess_with_profile_switch( stixify_reprocess_job, fake_stixifier_processor, stixifier_profile @@ -219,7 +275,8 @@ def test_process_post_reprocess_with_profile_switch( mock_stixify_processor_cls.return_value = fake_stixifier_processor process_post.si(stixify_reprocess_job.id).delay() stixify_reprocess_job.file.refresh_from_db() - fake_stixifier_processor.process.assert_called_once() + fake_stixifier_processor.file2txt.assert_called_once() + fake_stixifier_processor.txt2stix.assert_called_once() assert str(stixify_reprocess_job.file.profile_id) == str(new_profile.pk) mock_create_embedding.assert_called_once() @@ -274,11 +331,13 @@ def test_process_post__creates_embedding( @pytest.mark.django_db def test_process_post_full(stixify_job): - process_post.si(stixify_job.id).delay() - file = models.File.objects.get(pk=stixify_job.file_id) - stixify_job.refresh_from_db() - assert stixify_job.error == None, stixify_job.error - assert tuple(file.archived_pdf.read(4)) == (0x25, 0x50, 0x44, 0x46) + with patch("stixify.worker.pdf_converter.make_conversion") as mock_convert_pdf: + mock_convert_pdf.side_effect = lambda input_path, output_path: output_path.write_bytes(b"%PDF-1.4") + process_post.si(stixify_job.id).delay() + file = models.File.objects.get(pk=stixify_job.file_id) + stixify_job.refresh_from_db() + assert stixify_job.error == None, stixify_job.error + assert tuple(file.archived_pdf.read(4)) == (0x25, 0x50, 0x44, 0x46) @pytest.mark.django_db