Archive export and import with email and frontend notifications.

New NotificationCenter component.

Author: TontonSancho

colander/core/admin.py

@@ -2,6 +2,7 @@
 
 from colander.core.models import (
     Actor,
+    ArchiveExport,
     Artifact,
     ArtifactType,
     BackendCredentials,
@@ -15,12 +16,14 @@
     EntityRelation,
     Event,
     EventType,
+    FeedTemplate,
     Observable,
     ObservableRelation,
     ObservableType,
     PiRogueExperiment,
+    SubGraph,
     Threat,
-    UploadRequest, SubGraph, FeedTemplate,
+    UploadRequest, NotificationMessage,
 )
 
 
@@ -149,3 +152,12 @@ class SubGraphAdmin(admin.ModelAdmin):
 class FeedTemplateAdmin(admin.ModelAdmin):
     list_display = ('name', 'owner', 'visibility', 'in_error')
 admin.site.register(FeedTemplate, FeedTemplateAdmin)
+
+
+class ArchiveExportAdmin(admin.ModelAdmin):
+    list_display = ('case', 'type', 'requested_at', 'done_at', 'file')
+admin.site.register(ArchiveExport, ArchiveExportAdmin)
+
+class NotificationMessageAdmin(admin.ModelAdmin):
+    list_display = ('recipient', 'type', 'requested_at', 'processed_at', 'success')
+admin.site.register(NotificationMessage, NotificationMessageAdmin)

colander/core/api/views.py

@@ -290,5 +290,10 @@ class ApiDroppedFileViewSet(mixins.CreateModelMixin,
 
     def perform_create(self, serializer):
         inst = serializer.save(owner=self.request.user)
-        CaseContextConsumer.send_message_to_user_consumers(self.request.user, { 'msg': 'New dropped file' })
+        CaseContextConsumer.send_message_to_user_consumers(
+            self.request.user, {
+                'msg': 'A new drop is available',
+                'detail': inst.filename
+            }
+        )
         return inst

colander/core/archives/endpoints.py

@@ -0,0 +1,127 @@
+import json
+from functools import partial
+
+from django.contrib import messages
+from django.contrib.auth.decorators import login_required
+from django.db import transaction
+from django.http import JsonResponse, HttpResponseRedirect, HttpResponseForbidden, HttpResponse, \
+    StreamingHttpResponse
+from django.utils.http import content_disposition_header
+from rest_framework.exceptions import ValidationError
+
+from colander.core.archives.exporters import schedule_archive_export
+from colander.core.archives.serializers import model_by_super_types_str, serializers_by_model
+from colander.core.models import Case, ArchiveExport, Appendix
+
+
+@login_required
+def case_archive_request_view(request, pk):
+
+    case = Case.objects.get(pk=pk)
+
+    archive_export = ArchiveExport.objects.create(
+        case=case,
+        type=Appendix.ExportType.CASE,
+    )
+
+    #archives.schedule_archive_export(archive_export)
+    transaction.on_commit(partial(schedule_archive_export, archive_export))
+
+    messages.info(request, "Archive export requested. You will be notified when done.")
+
+    return HttpResponseRedirect(request.META.get('HTTP_REFERER'))
+
+
+@login_required
+def archive_takeout_view(request, pk):
+
+    archive_export = ArchiveExport.objects.get(pk=pk)
+    case = archive_export.case
+    if not case.can_contribute(request.user):
+        return HttpResponseForbidden()
+
+    if archive_export.is_pending:
+        response = HttpResponse(status=503)
+        response.headers["Retry-After"] = "120" # 2 minutes
+        return response
+
+    response = StreamingHttpResponse(archive_export.file, content_type='application/zip')
+    response['Content-Disposition'] = content_disposition_header(
+        True,
+        archive_export.filename
+    )
+
+    return response
+
+
+@login_required
+def archives_check_uuid_view(request, super_type, uuid):
+    if super_type not in model_by_super_types_str:
+        return JsonResponse({'message': f'{super_type} not supported'}, status=400)
+    model_class = model_by_super_types_str[super_type]
+    if model_class.objects.filter(pk=uuid).exists():
+        return JsonResponse({'message': 'uuid exists'}, status=200)
+    else:
+        return JsonResponse({'message': 'uuid does not exist'}, status=404)
+
+
+@login_required
+def archives_create_entity_view(request, super_type):
+    if super_type not in model_by_super_types_str:
+        return JsonResponse({'message': f'{super_type} not supported'}, status=400)
+
+    model_class = model_by_super_types_str[super_type]
+    serializer_class = serializers_by_model[model_class]
+
+    try:
+        payload = json.loads(request.body.decode('utf-8'))
+        serializer = serializer_class(data=payload, context={'request': request})
+        serializer.is_valid(raise_exception=True)
+        instance = serializer.save(owner=request.user)
+        return JsonResponse(serializer_class(instance).data)
+    except ValidationError as ve:
+        return JsonResponse({'message': ve.detail}, status=400)
+    except Exception as e:
+        return JsonResponse({'message': f'Unable to create entity: {e}'}, status=400)
+
+
+@login_required
+def archives_remap_entity_view(request, super_type, uuid):
+    if super_type not in model_by_super_types_str:
+        return JsonResponse({'message': f'{super_type} not supported'}, status=400)
+
+    model_class = model_by_super_types_str[super_type]
+    serializer_class = serializers_by_model[model_class]
+
+    try:
+        instance = model_class.objects.get(pk=uuid)
+        payload = json.loads(request.body.decode('utf-8'))
+        serializer = serializer_class(instance, data=payload, partial=True, context={'request': request})
+        serializer.is_valid(raise_exception=True)
+        instance = serializer.save()
+        return JsonResponse(serializer_class(instance).data)
+    except ValidationError as ve:
+        return JsonResponse({'message': ve.detail}, status=400)
+    except Exception as e:
+        return JsonResponse({'message': f'Unable to remap entity: {e}'}, status=400)
+
+
+@login_required
+def archives_attach_entity_view(request, super_type, uuid):
+    if super_type not in model_by_super_types_str:
+        return JsonResponse({'message': f'{super_type} not supported'}, status=400)
+
+    model_class = model_by_super_types_str[super_type]
+    serializer_class = serializers_by_model[model_class]
+
+    try:
+        instance = model_class.objects.get(pk=uuid)
+        payload = request.FILES
+        serializer = serializer_class(instance, data=payload, partial=True, context={'request': request})
+        serializer.is_valid(raise_exception=True)
+        instance = serializer.save()
+        return JsonResponse(serializer_class(instance).data)
+    except ValidationError as ve:
+        return JsonResponse({'message': ve.detail}, status=400)
+    except Exception as e:
+        return JsonResponse({'message': f'Unable to remap entity: {e}'}, status=400)

colander/core/archives/exporters/__init__.py

@@ -0,0 +1,80 @@
+
+
+import datetime
+
+from django.core.files.base import ContentFile
+from django.db.models.functions import Now
+from django_q.models import Schedule
+from django_q.tasks import async_task
+
+from colander.core.archives.exporters.case import CaseArchiveExporter
+from colander.core.models import ArchiveExport, Appendix
+from colander.core.notifications import notify_case_archive_done
+from colander.websocket.consumers import CaseContextConsumer
+
+MISSING_ARCHIVE_EXPORTER_SCHEDULE_TASK_NAME = 'missing_archive_exporter'
+
+
+def schedule_archive_export(archive_export:ArchiveExport):
+    hook=f'{_archive_export_end.__module__}.{_archive_export_end.__qualname__}'
+    async_task(_process_archive_export, str(archive_export.id), hook=hook)
+    _ensure_missing_archives_task_exist()
+
+
+def _process_archive_export(archive_export_id):
+    print(f'Processing archive export: {archive_export_id}')
+
+    archive_export = ArchiveExport.objects.get(pk=archive_export_id)
+    if archive_export.type == Appendix.ExportType.CASE:
+        cae = CaseArchiveExporter(archive_export.case)
+        ct = ContentFile(cae.export(pretty=True), 'dummy.zip')
+        archive_export.file = ct
+        archive_export.save()
+
+
+def _archive_export_end(task):
+    archive_export_id = task.args[0]
+
+    archive_export = ArchiveExport.objects.get(pk=archive_export_id)
+    if task.success:
+        archive_export.done_at = Now()
+        archive_export.save()
+        # Send notifications ...
+        CaseContextConsumer.send_message_to_user_consumers(archive_export.case.owner, {
+            'msg': 'A new archive is available',
+            'detail': archive_export.filename,
+        })
+        notify_case_archive_done(archive_export)
+
+    print("archive_export_done", archive_export_id, task.result)
+
+
+def _ensure_missing_archives_task_exist():
+    func=f'{_proceed_missing_archives.__module__}.{_proceed_missing_archives.__qualname__}'
+    print('scheduled func to call:', func)
+    if Schedule.objects.filter(func=func).exists(): return
+    print('creating scheduled func to call:', func)
+    Schedule.objects.create(func=func, schedule_type='H', name=MISSING_ARCHIVE_EXPORTER_SCHEDULE_TASK_NAME)
+
+
+def _proceed_missing_archives():
+    unprocessed_archive_exports = ArchiveExport.objects.filter(
+        done_at__isnull=True,
+        requested_at__lt=Now() - datetime.timedelta(hours=1)
+        ).order_by('-requested_at').all()
+
+    for archive_export_to_process in unprocessed_archive_exports:
+        schedule_archive_export(archive_export_to_process)
+
+    unprocessed_archive_exports_count = ArchiveExport.objects.filter(
+        done_at__isnull=True,
+        ).count()
+
+    if unprocessed_archive_exports_count > 0:
+        return
+
+    # unschedule 'missing_archive_exporter' task
+    func=f'{_proceed_missing_archives.__module__}.{_proceed_missing_archives.__qualname__}'
+    if Schedule.objects.filter(func=func).exists():
+        print('unscheduling func:', func)
+        Schedule.objects.get(func=func).delete()

colander/core/archives/exporters/case.py

@@ -0,0 +1,87 @@
+import datetime
+import io
+import json
+import zipfile
+from django.core.serializers.json import DjangoJSONEncoder
+
+from colander.core.archives.exporters.utils import Buffer
+from colander.core.archives.serializers import serializers_by_model
+from colander.core.models import Case, SubGraph
+
+
+class CaseArchiveExporter:
+    _case: Case
+
+    def __init__(self, case: Case):
+        self._case = case
+        self._entities = case.entities
+        self._entities.extend(SubGraph.objects.filter(case=self._case).all())
+        self._relations = case.relations
+        print('all entities', self._entities)
+
+    def export(self, pretty:bool=False):
+        buffer = io.BytesIO()
+        for buf in self.async_export(pretty=pretty):
+            buffer.write(buf)
+        buffer.seek(0)
+        return buffer.getvalue()
+
+    def async_export(self, pretty:bool=False):
+        manifest = {
+            'export_date': datetime.datetime.now(),
+            'files': {}
+        }
+        buffer = Buffer()
+        with zipfile.ZipFile(buffer, mode='w', compression=zipfile.ZIP_DEFLATED) as archive:
+            serializer = serializers_by_model.get(Case, None)
+            json_content = json.dumps( serializer(self._case, many=False).data, cls=DjangoJSONEncoder, indent=2 if pretty else None )
+
+            fname = "data.json"
+            archive.writestr("data.json", json_content)
+            yield buffer.take()
+            manifest['files'][fname] = 'md5sum'
+
+            for entity in self._entities:
+                entity_model = entity.__class__
+                serializer = serializers_by_model.get(entity_model, None)
+                if serializer:
+                    data = serializer(entity, many=False).data
+                    json_content = json.dumps( data, cls=DjangoJSONEncoder, indent=2 if pretty else None )
+                    fname = f"{str(entity_model.__name__)}/{str(entity.id)}/data.json"
+                    archive.writestr(fname, json_content)
+                    yield buffer.take()
+                    manifest['files'][fname] = 'md5sum'
+
+                    if hasattr(entity, 'thumbnail'):
+                        if entity.thumbnail:
+                            thumbnail_data = entity.thumbnail.file.read()
+                            fname = f"{str(entity_model.__name__)}/{str(entity.id)}/thumbnail.png"
+                            archive.writestr(fname, thumbnail_data)
+                            yield buffer.take()
+                            manifest['files'][fname] = 'md5sum'
+
+                    if hasattr(entity, 'file'):
+                        if entity.file:
+                            file_data = entity.file.file.read()
+                            fname = f"{str(entity_model.__name__)}/{str(entity.id)}/file{str(entity.extension)}"
+                            archive.writestr(fname, file_data)
+                            yield buffer.take()
+                            manifest['files'][fname] = 'md5sum'
+
+            for relation in self._relations:
+                if relation.immutable: continue
+                entity_model = relation.__class__
+                serializer = serializers_by_model.get(entity_model, None)
+                if serializer:
+                    data = serializer(relation, many=False).data
+                    json_content = json.dumps( data, cls=DjangoJSONEncoder, indent=2 if pretty else None )
+                    fname = f"{str(entity_model.__name__)}/{str(relation.id)}/data.json"
+                    archive.writestr(fname, json_content)
+                    yield buffer.take()
+                    manifest['files'][fname] = 'md5sum'
+
+            json_content = json.dumps( manifest, cls=DjangoJSONEncoder, indent=2 if pretty else None )
+            archive.writestr("manifest.json", json_content)
+            yield buffer.take()
+
+        yield buffer.end()

colander/core/archives/exporters/utils.py

@@ -0,0 +1,21 @@
+class Buffer:
+    # A lightweight buffer to use with zip library
+    def __init__(self):
+        self.buf = bytearray()
+
+    def write(self, data):
+        self.buf.extend(data)
+        return len(data)
+
+    def flush(self):
+        pass
+
+    def take(self):
+        buf = self.buf
+        self.buf = bytearray()
+        return bytes(buf)
+
+    def end(self):
+        buf = self.buf
+        self.buf = None
+        return bytes(buf)