Enhance upload behavior and validate labels and sources

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(

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,33 +146,28 @@ 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
                 file.save(update_fields=["profile"])
             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(

tests/src/test_tasks.py

@@ -183,14 +183,70 @@ 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()
         mock_convert_pdf.assert_not_called()
         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