From 26b504611fe5344c54ccc7a6255bc2df347784b1 Mon Sep 17 00:00:00 2001 From: lullah Date: Mon, 10 Aug 2026 12:45:05 +0100 Subject: [PATCH] allow reprocessing multiple files/reports at once #428 --- .env.example | 1 + .env.markdown | 2 + stixify/settings.py | 1 + .../0027_file_added_file_updated.py | 48 +++++ .../0028_rename_reprocess_job_type.py | 36 ++++ stixify/web/models.py | 10 +- stixify/web/serializers.py | 46 ++++- stixify/web/views.py | 134 +++++++++++++- stixify/worker/tasks.py | 155 ++++++++++++++-- tests/src/test_models.py | 13 +- tests/src/test_serializers.py | 40 ++++- tests/src/test_tasks.py | 169 +++++++++++++++++- tests/src/views/test_file_view.py | 163 +++++++++++++++++ 13 files changed, 793 insertions(+), 25 deletions(-) create mode 100644 stixify/web/migrations/0027_file_added_file_updated.py create mode 100644 stixify/web/migrations/0028_rename_reprocess_job_type.py diff --git a/.env.example b/.env.example index db7c9df..7133da7 100644 --- a/.env.example +++ b/.env.example @@ -15,6 +15,7 @@ CELERY_BROKER_CONNECTION_RETRY_ON_STARTUP= # stixify settings MAX_PAGE_SIZE= DEFAULT_PAGE_SIZE= +REPROCESS_MAX_FAILED_PROCESSES=10 # stix2arango settings ARANGODB_HOST_URL= ARANGODB_USERNAME= diff --git a/.env.markdown b/.env.markdown index c4defbc..a420ca8 100644 --- a/.env.markdown +++ b/.env.markdown @@ -36,6 +36,8 @@ These define how the API behaves. * This is the maximum number of results the API will ever return before pagination * `DEFAULT_PAGE_SIZE`: `50` * The default page size of result returned by the API +* `REPROCESS_MAX_FAILED_PROCESSES`: default `10` + * Stop a detached File reprocessing job before starting another File after this many failures. ## ArangoDB settings diff --git a/stixify/settings.py b/stixify/settings.py index 7476d63..2a392a4 100644 --- a/stixify/settings.py +++ b/stixify/settings.py @@ -280,3 +280,4 @@ CLASSIFIER_MODEL_PATH = os.getenv("CLASSIFIER_MODEL_PATH", os.path.join(BASE_DIR, "classifier_hdbscan.joblib")) CLASSIFIER_CONCURRENCY = int(os.getenv("CLASSIFIER_CONCURRENCY", 12)) CREATE_EMBEDDING_INCLUDE_NON_INCIDENT = bool(os.getenv("CREATE_EMBEDDING_INCLUDE_NON_INCIDENT", False)) +REPROCESS_MAX_FAILED_PROCESSES = int(os.getenv("REPROCESS_MAX_FAILED_PROCESSES", 10)) diff --git a/stixify/web/migrations/0027_file_added_file_updated.py b/stixify/web/migrations/0027_file_added_file_updated.py new file mode 100644 index 0000000..be64043 --- /dev/null +++ b/stixify/web/migrations/0027_file_added_file_updated.py @@ -0,0 +1,48 @@ +from django.db import migrations, models + + +def populate_file_timestamps(apps, schema_editor): + File = apps.get_model("stixify_core", "File") + Job = apps.get_model("stixify_core", "Job") + first_job_times = dict( + Job.objects.filter(file_id__isnull=False) + .order_by("file_id", "run_datetime", "id") + .distinct("file_id") + .values_list("file_id", "run_datetime") + ) + files = list(File.objects.all()) + for file in files: + timestamp = first_job_times.get(file.pk, file.modified) + file.added = timestamp + file.updated = timestamp + File.objects.bulk_update(files, ["added", "updated"]) + + +class Migration(migrations.Migration): + dependencies = [ + ("stixify_core", "0026_file_pap_level"), + ] + + operations = [ + migrations.AddField( + model_name="file", + name="added", + field=models.DateTimeField(null=True), + ), + migrations.AddField( + model_name="file", + name="updated", + field=models.DateTimeField(null=True), + ), + migrations.RunPython(populate_file_timestamps, migrations.RunPython.noop), + migrations.AlterField( + model_name="file", + name="added", + field=models.DateTimeField(auto_now_add=True), + ), + migrations.AlterField( + model_name="file", + name="updated", + field=models.DateTimeField(auto_now=True), + ), + ] diff --git a/stixify/web/migrations/0028_rename_reprocess_job_type.py b/stixify/web/migrations/0028_rename_reprocess_job_type.py new file mode 100644 index 0000000..6350c90 --- /dev/null +++ b/stixify/web/migrations/0028_rename_reprocess_job_type.py @@ -0,0 +1,36 @@ +from django.db import migrations, models + + +def rename_reprocess_job_type(apps, schema_editor): + Job = apps.get_model("stixify_core", "Job") + Job.objects.filter(type="reprocess-posts").update(type="reprocess-files") + + +def restore_reprocess_job_type(apps, schema_editor): + Job = apps.get_model("stixify_core", "Job") + Job.objects.filter(type="reprocess-files").update(type="reprocess-posts") + + +class Migration(migrations.Migration): + dependencies = [ + ("stixify_core", "0027_file_added_file_updated"), + ] + + operations = [ + migrations.RunPython(rename_reprocess_job_type, restore_reprocess_job_type), + migrations.AlterField( + model_name="job", + name="type", + field=models.CharField( + choices=[ + ("import-file", "Import File"), + ("reprocess-files", "Reprocess Files"), + ("sync-knowledgebase", "Sync Knowledgebase"), + ("build-clusters", "Build Clusters"), + ("build-embeddings", "Build Embeddings"), + ], + default="import-file", + max_length=64, + ), + ), + ] diff --git a/stixify/web/models.py b/stixify/web/models.py index b9e7616..e06c2b1 100644 --- a/stixify/web/models.py +++ b/stixify/web/models.py @@ -118,6 +118,8 @@ def validate_file(file: InMemoryUploadedFile, mode: str): class File(CommonSTIXProps): id = models.UUIDField(unique=True, max_length=64, primary_key=True, default=uuid.uuid4) + added = models.DateTimeField(auto_now_add=True) + updated = models.DateTimeField(auto_now=True) file = models.FileField(max_length=1024, upload_to=upload_to_func) identity = models.ForeignKey(Identity, on_delete=models.CASCADE, default=None) profile = models.ForeignKey(Profile, on_delete=models.PROTECT) @@ -174,6 +176,12 @@ def report_id(self, value): def clean(self) -> None: validate_file(self.file, self.mode) return super().clean() + + def save(self, *args, **kwargs): + update_fields = kwargs.get("update_fields") + if update_fields is not None: + kwargs["update_fields"] = set(update_fields) | {"updated"} + return super().save(*args, **kwargs) def __str__(self) -> str: return f"File(id={self.id})" @@ -298,7 +306,7 @@ class JobState(models.TextChoices): class JobType(models.TextChoices): IMPORT_FILE = "import-file" - REPROCESS_POSTS = "reprocess-posts" + REPROCESS_FILES = "reprocess-files" SYNC_KNOWLEDGEBASE = "sync-knowledgebase" BUILD_CLUSTERS = "build-clusters" BUILD_EMBEDDINGS = "build-embeddings" diff --git a/stixify/web/serializers.py b/stixify/web/serializers.py index 30a2d2e..bd5116c 100644 --- a/stixify/web/serializers.py +++ b/stixify/web/serializers.py @@ -144,7 +144,7 @@ class FileSerializer(serializers.ModelSerializer): class Meta: model = File exclude = ['profile', "markdown_file", "txt2stix_data", "pdf_file", "identity", "embedding"] - read_only_fields = [] + read_only_fields = ["added", "updated"] def validate(self, attrs): return super().validate(attrs) @@ -240,6 +240,50 @@ def validate(self, attrs): return super().validate(attrs) +class ReprocessFilesSerializer(ReprocessSingleFileSerializer): + file_ids = serializers.ListField( + child=serializers.UUIDField(), + required=False, + help_text="IDs of the Files to reprocess.", + ) + identity_id = IdentityIDField( + required=False, + help_text="Reprocess Files owned by this Identity.", + ) + added_after = serializers.DateTimeField( + required=False, + help_text="Only reprocess Files added at or after this time.", + ) + added_before = serializers.DateTimeField( + required=False, + help_text="Only reprocess Files added at or before this time.", + ) + + def validate(self, attrs): + attrs = super().validate(attrs) + if not attrs.get("file_ids") and not attrs.get("identity_id"): + raise serializers.ValidationError( + { + "non_field_errors": [ + "At least one of file_ids or identity_id must be provided" + ] + } + ) + if ( + attrs.get("added_after") + and attrs.get("added_before") + and attrs["added_after"] > attrs["added_before"] + ): + raise serializers.ValidationError( + { + "non_field_errors": [ + "added_after cannot be later than added_before" + ] + } + ) + return attrs + + class AttackNavigatorDomainSerializer(JSONSchemaSerializer): json_schema = { "$schema": "http://json-schema.org/draft-07/schema#", diff --git a/stixify/web/views.py b/stixify/web/views.py index 8aaf5dc..9b1d3af 100644 --- a/stixify/web/views.py +++ b/stixify/web/views.py @@ -63,6 +63,7 @@ HealthCheckSerializer, ImageSerializer, JobSerializer, + ReprocessFilesSerializer, ReprocessSingleFileSerializer, ) from .topics import SimilarFileSerializer @@ -202,6 +203,7 @@ def _get_schema_response(self, request): ), ), reprocess=extend_schema( + operation_id="v1_files_reprocess_single", summary="Reprocess an uploaded File", description=textwrap.dedent( """ @@ -221,6 +223,24 @@ def _get_schema_response(self, request): }, request=ReprocessSingleFileSerializer, ), + reprocess_files=extend_schema( + operation_id="v1_files_reprocess_multiple", + summary="Reprocess uploaded Files", + description=textwrap.dedent( + """ + Reprocess multiple Files selected by `file_ids`, `identity_id`, or both. + + `added_after` and `added_before` can be used to restrict the resolved Files + by their immutable `added` timestamp. The response is one detached Job whose + `extra.file_ids` contains the complete resolved set of Files. + """ + ), + responses={ + 201: JobSerializer, + 400: DEFAULT_400_ERROR, + }, + request=ReprocessFilesSerializer, + ), ) class FileView( mixins.CreateModelMixin, @@ -536,7 +556,110 @@ def reprocess(self, request, file_id=None, **kwargs): raise exceptions.ValidationError( {"error": "Cannot skip extraction on unprocessed file"} ) - job = tasks.create_reprocessing_job(file_obj, s.validated_data) + job = tasks.create_reprocessing_job( + [file_obj.id], self._reprocess_options(s.validated_data) + ) + return Response( + JobSerializer(job, context={"request": request}).data, + status=status.HTTP_201_CREATED, + ) + + @staticmethod + def _reprocess_options(validated_data): + options = {} + for key in ( + "identity_id", + "profile_id", + "skip_extraction", + "added_after", + "added_before", + ): + if key not in validated_data: + continue + value = validated_data[key] + if hasattr(value, "isoformat"): + value = value.isoformat() + elif value is not None: + value = str(value) + options[key] = value + return options + + @decorators.action( + methods=["PATCH"], detail=False, url_path="reprocess" + ) + def reprocess_files(self, request, **kwargs): + serializer = ReprocessFilesSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + data = serializer.validated_data + + requested_ids = [str(file_id) for file_id in data.get("file_ids", [])] + existing_ids = { + str(file_id) + for file_id in File.objects.filter(pk__in=requested_ids).values_list( + "pk", flat=True + ) + } + missing_ids = [ + file_id for file_id in requested_ids if file_id not in existing_ids + ] + if missing_ids: + raise exceptions.ValidationError( + {"file_ids": [f"Unknown File IDs: {', '.join(missing_ids)}"]} + ) + + files = File.objects.all() + if data.get("added_after"): + files = files.filter(added__gte=data["added_after"]) + if data.get("added_before"): + files = files.filter(added__lte=data["added_before"]) + + eligible_explicit_ids = { + str(file_id) + for file_id in files.filter(pk__in=requested_ids).values_list("pk", flat=True) + } + resolved_ids = [] + seen = set() + for file_id in requested_ids: + if file_id in eligible_explicit_ids and file_id not in seen: + resolved_ids.append(file_id) + seen.add(file_id) + + if identity_id := data.get("identity_id"): + identity_file_ids = files.filter(identity_id=identity_id).order_by( + "added", "id" + ).values_list("pk", flat=True) + for file_id in identity_file_ids: + file_id = str(file_id) + if file_id not in seen: + resolved_ids.append(file_id) + seen.add(file_id) + + if not resolved_ids: + raise exceptions.ValidationError( + {"file_ids": ["No Files matched the supplied selectors"]} + ) + + if data["skip_extraction"]: + unprocessed_ids = [ + str(file_id) + for file_id, extraction_data in File.objects.filter( + pk__in=resolved_ids + ).values_list("pk", "txt2stix_data") + if not extraction_data + ] + if unprocessed_ids: + raise exceptions.ValidationError( + { + "file_ids": [ + "Cannot skip extraction for unprocessed Files: " + + ", ".join(unprocessed_ids) + ] + } + ) + + job = tasks.create_reprocessing_job( + resolved_ids, self._reprocess_options(data) + ) return Response( JobSerializer(job, context={"request": request}).data, status=status.HTTP_201_CREATED, @@ -589,11 +712,18 @@ def get_queryset(self): return Job.objects.all() class filterset_class(FilterSet): - file_id = Filter("file_id", label="Filter Jobs by File `id`") + file_id = filters.UUIDFilter( + method="filter_file_id", label="Filter Jobs by File `id`" + ) state = filters.BaseCSVFilter( help_text="Filter Jobs by their state.", lookup_expr="in" ) + def filter_file_id(self, queryset, name, value): + return queryset.filter( + Q(file_id=value) | Q(extra__file_ids__contains=[str(value)]) + ) + @extend_schema_view( list=extend_schema( diff --git a/stixify/worker/tasks.py b/stixify/worker/tasks.py index fdff9a4..6d1424c 100644 --- a/stixify/worker/tasks.py +++ b/stixify/worker/tasks.py @@ -9,8 +9,9 @@ from txt2stix import txt2stixBundler from stixify.web.models import Job, File from stixify.web import models -from celery import shared_task +from celery import chain, shared_task from dogesec_commons.stixifier.stixifier import StixifyProcessor, ReportProperties +from dogesec_commons.stixifier.models import Profile from stixify.web.values.statistics import build_data_and_add_to_cache from django.core.files.uploadedfile import InMemoryUploadedFile @@ -66,16 +67,41 @@ def release_upload_lock(job_id): def new_task(job: Job): - (process_post.s(job.id) | job_completed_with_error.si(job.id)).apply_async( + if job.type == models.JobType.REPROCESS_FILES and (job.extra or {}).get("file_ids"): + task = chain( + *[ + process_post.si(job.id, file_id) + for file_id in job.extra["file_ids"] + ], + job_completed_with_error.si(job.id), + ) + else: + task = process_post.s(job.id) | job_completed_with_error.si(job.id) + task.apply_async( countdown=POLL_INTERVAL, root_id=str(job.id), task_id=str(job.id) ) -def create_reprocessing_job(file: File, options: dict = None): - options = options or {} - job = models.Job.objects.create( +def create_reprocessing_job(file_ids, options: dict = None): + file_ids = [str(file_id) for file_id in file_ids] + options = dict(options or {}) + options.update( + file_ids=file_ids, + progress=dict( + total_items=len(file_ids), + processed_items=0, + failed_processes=0, + unprocessed_items=len(file_ids), + current_file_id=None, + current_index=None, + stopped_early=False, + stop_reason=None, + errors=[], + ), + ) + job = models.Job.objects.create( id=uuid.uuid4(), - type=models.JobType.REPROCESS_POSTS, - file=file, + type=models.JobType.REPROCESS_FILES, + file=None, state=models.JobState.PENDING, extra=options, ) @@ -84,7 +110,7 @@ def create_reprocessing_job(file: File, options: dict = None): def _process_file(processor, job, file): skip_extraction = bool((job.extra or {}).get("skip_extraction")) - is_reprocess = job.type == models.JobType.REPROCESS_POSTS + is_reprocess = job.type == models.JobType.REPROCESS_FILES if is_reprocess and skip_extraction: processor.output_md = file.markdown_file.open().read().decode() @@ -101,16 +127,79 @@ def _process_file(processor, job, file): processor.write_bundle(processor.bundler) +def _object_value_backup(file_id): + return list( + models.ObjectValue.objects.filter(file_id=file_id).values( + "id", + "stix_id", + "type", + "knowledgebase", + "values", + "created", + "modified", + "is_dupe", + ) + ) + + +def _restore_object_values(file_id, backup): + models.ObjectValue.objects.filter(file_id=file_id).delete() + models.ObjectValue.objects.bulk_create( + [models.ObjectValue(file_id=file_id, **values) for values in backup] + ) + + +def _update_reprocess_progress(job, file_id, error=None): + progress = job.extra["progress"] + if error is None: + progress["processed_items"] += 1 + else: + progress["failed_processes"] += 1 + progress["errors"].append( + {"file_id": str(file_id), "message": error} + ) + if progress["failed_processes"] >= settings.REPROCESS_MAX_FAILED_PROCESSES: + progress["stopped_early"] = True + progress["stop_reason"] = "failure_limit_reached" + progress["unprocessed_items"] = max( + 0, + progress["total_items"] + - progress["processed_items"] + - progress["failed_processes"], + ) + + @shared_task -def process_post(job_id, *args): +def process_post(job_id, file_id=None, *args): job = Job.objects.get(id=job_id) - file = job.file + detached_reprocess = ( + job.type == models.JobType.REPROCESS_FILES and file_id is not None + ) + if detached_reprocess: + progress = job.extra["progress"] + if progress["failed_processes"] >= settings.REPROCESS_MAX_FAILED_PROCESSES: + return job_id + progress["current_file_id"] = str(file_id) + progress["current_index"] = ( + progress["processed_items"] + progress["failed_processes"] + ) + file = None + else: + file = job.file + object_values_backup = None try: + if detached_reprocess: + file = File.objects.get(pk=file_id) job.state = models.JobState.PROCESSING job.save() + processing_profile = file.profile + if job.type == models.JobType.REPROCESS_FILES and (job.extra or {}).get( + "profile_id" + ): + processing_profile = Profile.objects.get(pk=job.extra["profile_id"]) processor = StixifyProcessor( file.process_file, - job.profile, + processing_profile, job_id=job.id, file2txt_mode=file.process_mode, report_id=file.id, @@ -118,7 +207,7 @@ def process_post(job_id, *args): external_refs = [ dict( source_name="stixify_profile_id", - external_id=str(job.profile.id), + external_id=str(processing_profile.id), ) ] for source in file.sources or []: @@ -147,9 +236,12 @@ def process_post(job_id, *args): report_prop=report_props, extra=dict(_stixify_file_id=str(file.id)) ) - models.ObjectValue.objects.filter(file_id=file.id).delete() _process_file(processor, job, file) + if job.type == models.JobType.REPROCESS_FILES: + object_values_backup = _object_value_backup(file.id) + models.ObjectValue.objects.filter(file_id=file.id).delete() + acquire_upload_lock(job.id) try: logging.info(f"uploading {processor.task_name} to arangodb via stix2arango") @@ -173,7 +265,7 @@ def process_post(job_id, *args): models.FileImage.objects.create( report=file, file=DjangoFile(image, image.name), name=image.name ) - if job.profile.generate_pdf: + if processing_profile.generate_pdf: converted_file_path = processor.tmpdir / "converted_pdf.pdf" pdf_converter.make_conversion(processor.filename, converted_file_path) file.pdf_file.save( @@ -182,11 +274,23 @@ def process_post(job_id, *args): file.save(update_fields=['markdown_file', 'pdf_file']) except Exception as e: error = str(e) - job.error = "failed to process report" + job.error = "failed to process file" if error: job.error += f": {error}" + if object_values_backup is not None and file is not None: + try: + _restore_object_values(file.id, object_values_backup) + except Exception: + logging.exception( + "failed to restore ObjectValue data for File %s", file.id + ) + if detached_reprocess: + _update_reprocess_progress(job, file_id, job.error) logging.error(job.error) logging.exception(e) + else: + if detached_reprocess: + _update_reprocess_progress(job, file_id) job.save() return job_id @@ -195,11 +299,30 @@ def process_post(job_id, *args): def job_completed_with_error(job_id): job = Job.objects.get(pk=job_id) state = models.JobState.COMPLETED + if job.type == models.JobType.REPROCESS_FILES and (job.extra or {}).get("progress"): + progress = job.extra["progress"] + progress["current_file_id"] = None + progress["current_index"] = None + progress["unprocessed_items"] = max( + 0, + progress["total_items"] + - progress["processed_items"] + - progress["failed_processes"], + ) + if progress["failed_processes"]: + job.error = ( + f"failed to reprocess {progress['failed_processes']} file(s)" + ) if job.error: state = models.JobState.FAILED if job.type == models.JobType.IMPORT_FILE: job.file and job.file.delete() - Job.objects.filter(pk=job_id).update(state=state, completion_time=datetime.now(UTC)) + Job.objects.filter(pk=job_id).update( + state=state, + error=job.error, + extra=job.extra, + completion_time=datetime.now(UTC), + ) from celery import signals diff --git a/tests/src/test_models.py b/tests/src/test_models.py index 43edb6b..741d7f7 100644 --- a/tests/src/test_models.py +++ b/tests/src/test_models.py @@ -4,4 +4,15 @@ def test_upload_to_func(db, stixify_file): image = models.FileImage.objects.create(report=stixify_file) assert models.upload_to_func(stixify_file, "ade.pdf") == "identity--c5f27ca2-a580-4fee-9bb9-753e2b563a30/report--dcbeb240-8dd6-4892-8e9e-7b6bda30e454/dcbeb240-8dd6-4892-8e9e-7b6bda30e454_ade.pdf" - assert models.upload_to_func(image, "ade.png") == "identity--c5f27ca2-a580-4fee-9bb9-753e2b563a30/report--dcbeb240-8dd6-4892-8e9e-7b6bda30e454/dcbeb240-8dd6-4892-8e9e-7b6bda30e454_ade.png" \ No newline at end of file + assert models.upload_to_func(image, "ade.png") == "identity--c5f27ca2-a580-4fee-9bb9-753e2b563a30/report--dcbeb240-8dd6-4892-8e9e-7b6bda30e454/dcbeb240-8dd6-4892-8e9e-7b6bda30e454_ade.png" + + +def test_file_added_is_immutable_and_updated_changes(db, stixify_file): + original_added = stixify_file.added + original_updated = stixify_file.updated + stixify_file.name = "Updated name" + stixify_file.save(update_fields=["name"]) + stixify_file.refresh_from_db() + + assert stixify_file.added == original_added + assert stixify_file.updated >= original_updated diff --git a/tests/src/test_serializers.py b/tests/src/test_serializers.py index 64ecd37..659ad6b 100644 --- a/tests/src/test_serializers.py +++ b/tests/src/test_serializers.py @@ -1,6 +1,10 @@ import pytest from django.core.files.base import ContentFile -from stixify.web.serializers import FileSerializer, FilePatchSerializer +from stixify.web.serializers import ( + FilePatchSerializer, + FileSerializer, + ReprocessFilesSerializer, +) @pytest.mark.django_db @@ -199,3 +203,37 @@ def test_create_with_invalid_sources(self, stixify_job): serializer = FileSerializer(data=data) assert not serializer.is_valid() assert "sources" in serializer.errors + + +@pytest.mark.django_db +class TestReprocessFilesSerializer: + def test_requires_file_ids_or_identity_id(self): + serializer = ReprocessFilesSerializer(data={"skip_extraction": True}) + + assert not serializer.is_valid() + assert "non_field_errors" in serializer.errors + + def test_rejects_reversed_added_range(self, stixify_file): + serializer = ReprocessFilesSerializer( + data={ + "file_ids": [str(stixify_file.id)], + "skip_extraction": True, + "added_after": "2026-08-02T00:00:00Z", + "added_before": "2026-08-01T00:00:00Z", + } + ) + + assert not serializer.is_valid() + assert "non_field_errors" in serializer.errors + + def test_accepts_file_ids_and_added_range(self, stixify_file): + serializer = ReprocessFilesSerializer( + data={ + "file_ids": [str(stixify_file.id)], + "skip_extraction": True, + "added_after": "2026-08-01T00:00:00Z", + "added_before": "2026-08-02T00:00:00Z", + } + ) + + assert serializer.is_valid(), serializer.errors diff --git a/tests/src/test_tasks.py b/tests/src/test_tasks.py index bb5f6d7..7f226e3 100644 --- a/tests/src/test_tasks.py +++ b/tests/src/test_tasks.py @@ -8,6 +8,7 @@ from dogesec_commons.stixifier.stixifier import StixifyProcessor from dogesec_commons.stixifier.models import Profile from django.core.files.base import ContentFile +from django.test import override_settings from txt2stix.txt2stix import Txt2StixData from stixify.worker import tasks @@ -31,6 +32,41 @@ def test_new_task(stixify_job): mock_job_completed_with_error.assert_called_once_with(stixify_job.id) +@pytest.mark.django_db +def test_new_task_detached_reprocesses_each_file(stixify_file): + file_ids = [str(stixify_file.id), str(stixify_file.id)] + job = models.Job.objects.create( + type=models.JobType.REPROCESS_FILES, + extra={ + "file_ids": file_ids, + "progress": { + "total_items": 2, + "processed_items": 0, + "failed_processes": 0, + "unprocessed_items": 2, + "current_file_id": None, + "current_index": None, + "stopped_early": False, + "stop_reason": None, + "errors": [], + }, + }, + ) + with ( + patch("stixify.worker.tasks.process_post.run") as mock_process_file, + patch( + "stixify.worker.tasks.job_completed_with_error.run" + ) as mock_completed, + ): + new_task(job) + + assert mock_process_file.call_args_list == [ + call(job.id, file_ids[0]), + call(job.id, file_ids[1]), + ] + mock_completed.assert_called_once_with(job.id) + + @pytest.mark.django_db def test_process_post_job__fails(stixify_job): with ( @@ -40,12 +76,12 @@ def test_process_post_job__fails(stixify_job): ): process_post.si(stixify_job.id).delay() stixify_job.refresh_from_db() - assert stixify_job.error == "failed to process report" + assert stixify_job.error == "failed to process file" mock_stixify_processor_cls.side_effect = ValueError("some error") process_post.si(stixify_job.id).delay() stixify_job.refresh_from_db() - assert stixify_job.error == "failed to process report: some error" + assert stixify_job.error == "failed to process file: some error" @pytest.fixture @@ -63,7 +99,7 @@ def fake_stixifier_processor(tmpdir): @pytest.fixture def stixify_reprocess_job(stixify_job): - stixify_job.type = models.JobType.REPROCESS_POSTS + stixify_job.type = models.JobType.REPROCESS_FILES stixify_job.extra = {} stixify_job.save(update_fields=["type", "extra"]) stixify_job.file.set_txt2stix_data(fake_txt2stix_data()) @@ -244,6 +280,7 @@ def test_process_post_concurrent_uploads_limited( from stixify.worker.tasks import acquire_upload_lock with pytest.raises(TimeoutError): acquire_upload_lock(stixify_reprocess_job.id, wait_timeout=0.1) + cache.clear() @@ -278,6 +315,7 @@ def test_process_post_reprocess_with_profile_switch( 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) + assert mock_stixify_processor_cls.call_args.args[1] == new_profile mock_create_embedding.assert_called_once() @@ -362,3 +400,128 @@ def test_job_completed_with_error__success(stixify_job): assert stixify_job.file.pk == file_id assert stixify_job.state == models.JobState.COMPLETED assert stixify_job.completion_time != None + + +def detached_reprocess_job(file_ids, **options): + progress = { + "total_items": len(file_ids), + "processed_items": 0, + "failed_processes": 0, + "unprocessed_items": len(file_ids), + "current_file_id": None, + "current_index": None, + "stopped_early": False, + "stop_reason": None, + "errors": [], + } + return models.Job.objects.create( + type=models.JobType.REPROCESS_FILES, + extra={"file_ids": file_ids, "progress": progress, **options}, + ) + + +@pytest.mark.django_db +def test_detached_reprocess_updates_progress( + stixify_file, fake_stixifier_processor +): + job = detached_reprocess_job([str(stixify_file.id)], skip_extraction=False) + with ( + patch("stixify.worker.tasks.StixifyProcessor") as processor_class, + patch.object(models.File, "create_embedding"), + ): + processor_class.return_value = fake_stixifier_processor + process_post(job.id, stixify_file.id) + + job.refresh_from_db() + assert job.extra["progress"] == { + "total_items": 1, + "processed_items": 1, + "failed_processes": 0, + "unprocessed_items": 0, + "current_file_id": str(stixify_file.id), + "current_index": 0, + "stopped_early": False, + "stop_reason": None, + "errors": [], + } + + +@pytest.mark.django_db +@override_settings(REPROCESS_MAX_FAILED_PROCESSES=10) +def test_detached_reprocess_stops_after_ten_failures(stixify_file): + file_ids = [str(stixify_file.id)] * 11 + job = detached_reprocess_job(file_ids, skip_extraction=False) + + with patch( + "stixify.worker.tasks.StixifyProcessor", side_effect=ValueError("bad file") + ) as processor_class: + for file_id in file_ids: + process_post(job.id, file_id) + + job.refresh_from_db() + progress = job.extra["progress"] + assert processor_class.call_count == 10 + assert progress["processed_items"] == 0 + assert progress["failed_processes"] == 10 + assert progress["unprocessed_items"] == 1 + assert progress["stopped_early"] is True + assert progress["stop_reason"] == "failure_limit_reached" + assert len(progress["errors"]) == 10 + assert set(progress["errors"][0]) == {"file_id", "message"} + + job_completed_with_error(job.id) + job.refresh_from_db() + assert job.state == models.JobState.FAILED + assert job.error == "failed to reprocess 10 file(s)" + assert job.extra["progress"]["current_file_id"] is None + + +@pytest.mark.django_db +def test_detached_reprocess_is_failed_when_one_file_fails( + stixify_file, fake_stixifier_processor +): + file_ids = [str(stixify_file.id), str(stixify_file.id)] + job = detached_reprocess_job(file_ids, skip_extraction=False) + + with ( + patch( + "stixify.worker.tasks.StixifyProcessor", + side_effect=[ValueError("bad file"), fake_stixifier_processor], + ), + patch.object(models.File, "create_embedding"), + ): + for file_id in file_ids: + process_post(job.id, file_id) + job_completed_with_error(job.id) + + job.refresh_from_db() + assert job.state == models.JobState.FAILED + assert job.extra["progress"]["processed_items"] == 1 + assert job.extra["progress"]["failed_processes"] == 1 + assert job.extra["progress"]["unprocessed_items"] == 0 + + +@pytest.mark.django_db +def test_reprocess_restores_object_values_after_failure( + stixify_file, fake_stixifier_processor +): + original = models.ObjectValue.objects.create( + file=stixify_file, + stix_id="indicator--11111111-1111-4111-8111-111111111111", + type="indicator", + values={"name": "original"}, + ) + job = detached_reprocess_job([str(stixify_file.id)], skip_extraction=False) + fake_stixifier_processor.upload_to_arango.side_effect = RuntimeError("upload failed") + + with ( + patch("stixify.worker.tasks.StixifyProcessor") as processor_class, + patch.object(models.File, "create_embedding"), + ): + processor_class.return_value = fake_stixifier_processor + process_post(job.id, stixify_file.id) + + restored = models.ObjectValue.objects.get( + file=stixify_file, stix_id=original.stix_id + ) + assert restored.values == {"name": "original"} diff --git a/tests/src/views/test_file_view.py b/tests/src/views/test_file_view.py index 8ffe4c1..30fd816 100644 --- a/tests/src/views/test_file_view.py +++ b/tests/src/views/test_file_view.py @@ -10,6 +10,8 @@ from django.core.files.uploadedfile import SimpleUploadedFile from django.core.files.base import ContentFile import io +from datetime import timedelta +from django.utils import timezone from stixify.web.md_helper import MarkdownImageReplacer from tests.utils import Transport @@ -585,3 +587,164 @@ def test_file_pdf_with_pdf(client, stixify_file, api_schema): assert re.match(r'attachment; filename="dcbeb240-8dd6-4892-8e9e-7b6bda30e454_archived_*[\w]*.pdf"', resp.headers["Content-Disposition"]) assert resp.getvalue() == b"pdf content" api_schema['/api/v1/files/{file_id}/pdf/']['GET'].validate_response(Transport.get_st_response(resp)) + + +@pytest.mark.django_db +def test_reprocess_single_file_creates_detached_job( + client, stixify_file, api_schema +): + stixify_file.txt2stix_data = {"existing": True} + stixify_file.save(update_fields=["txt2stix_data"]) + + with patch("stixify.worker.tasks.new_task") as mock_new_task: + response = client.patch( + f"/api/v1/files/{stixify_file.id}/reprocess/", + data=json.dumps({"skip_extraction": True}), + content_type="application/json", + ) + + assert response.status_code == 201, response.content + job = models.Job.objects.get(pk=response.data["id"]) + assert job.file is None + assert job.type == models.JobType.REPROCESS_FILES + assert job.extra["file_ids"] == [str(stixify_file.id)] + assert job.extra["progress"]["total_items"] == 1 + mock_new_task.assert_called_once_with(job) + api_schema["/api/v1/files/{file_id}/reprocess/"]["PATCH"].validate_response( + Transport.get_st_response(response) + ) + + +@pytest.mark.django_db +def test_reprocess_files_creates_one_detached_job( + client, more_files, api_schema +): + requested_ids = [ + str(more_files[1].id), + str(more_files[0].id), + str(more_files[1].id), + ] + with patch("stixify.worker.tasks.new_task") as mock_new_task: + response = client.patch( + "/api/v1/files/reprocess/", + data=json.dumps( + {"file_ids": requested_ids, "skip_extraction": True} + ), + content_type="application/json", + ) + + assert response.status_code == 201, response.content + job = models.Job.objects.get(pk=response.data["id"]) + assert job.file is None + assert job.extra["file_ids"] == requested_ids[:2] + assert job.extra["progress"]["total_items"] == 2 + assert job.extra["progress"]["unprocessed_items"] == 2 + mock_new_task.assert_called_once_with(job) + api_schema["/api/v1/files/reprocess/"]["PATCH"].validate_response( + Transport.get_st_response(response) + ) + + +@pytest.mark.django_db +def test_reprocess_files_filters_identity_by_added_date(client, more_files, identity): + now = timezone.now() + models.File.objects.filter(pk=more_files[0].pk).update( + added=now - timedelta(days=2) + ) + models.File.objects.filter(pk=more_files[1].pk).update(added=now) + models.File.objects.filter(pk=more_files[2].pk).update( + added=now + timedelta(days=2) + ) + + with patch("stixify.worker.tasks.new_task"): + response = client.patch( + "/api/v1/files/reprocess/", + data=json.dumps( + { + "identity_id": str(identity.id), + "added_after": (now - timedelta(days=1)).isoformat(), + "added_before": (now + timedelta(days=1)).isoformat(), + "skip_extraction": True, + } + ), + content_type="application/json", + ) + + assert response.status_code == 201, response.content + assert response.data["extra"]["file_ids"] == [str(more_files[1].id)] + assert response.data["extra"]["added_after"] == ( + now - timedelta(days=1) + ).isoformat() + + +@pytest.mark.django_db +def test_reprocess_files_requires_selector(client): + response = client.patch( + "/api/v1/files/reprocess/", + data=json.dumps({"skip_extraction": True}), + content_type="application/json", + ) + assert response.status_code == 400 + assert "non_field_errors" in json.loads(response.content)["details"] + + +@pytest.mark.django_db +def test_jobs_file_id_filter_includes_detached_reprocessing_job( + client, stixify_file +): + attached = models.Job.objects.create(file=stixify_file) + detached = models.Job.objects.create( + type=models.JobType.REPROCESS_FILES, + extra={"file_ids": [str(stixify_file.id)]}, + ) + + response = client.get( + "/api/v1/jobs/", query_params={"file_id": str(stixify_file.id)} + ) + + assert response.status_code == 200, response.content + assert {str(job["id"]) for job in response.data["jobs"]} == { + str(attached.id), + str(detached.id), + } + + +@pytest.mark.django_db +def test_reprocess_files_rejects_unknown_file_ids(client): + unknown_file_id = str(uuid.uuid4()) + response = client.patch( + "/api/v1/files/reprocess/", + data=json.dumps( + {"file_ids": [unknown_file_id], "skip_extraction": True} + ), + content_type="application/json", + ) + + assert response.status_code == 400 + assert unknown_file_id in str( + json.loads(response.content)["details"]["file_ids"] + ) + + +@pytest.mark.django_db +def test_reprocess_files_validates_skip_extraction_for_entire_batch( + client, more_files +): + more_files[1].txt2stix_data = None + more_files[1].save(update_fields=["txt2stix_data"]) + + response = client.patch( + "/api/v1/files/reprocess/", + data=json.dumps( + { + "file_ids": [str(file.id) for file in more_files[:2]], + "skip_extraction": True, + } + ), + content_type="application/json", + ) + + assert response.status_code == 400 + assert str(more_files[1].id) in str( + json.loads(response.content)["details"]["file_ids"] + )