Added whitelist of file types in file upload (#1824)

* feat: added whitelist of file types in file upload

* test: added test for file upload whitelist

* feat: use mime type for files whitelisting

* test: use mime_type in test instead of extension

* fix: allow audio files

* fix: added all the allowed file types to whitelist

* fix: minor changes (implement review suggestions)

* Fix Typo

---------

Co-authored-by: Vignesh Hari <[email protected]>

care/facility/api/serializers/file_upload.py

@@ -1,6 +1,7 @@
+from django.conf import settings
 from django.utils.timezone import localtime, now
-from jsonschema import ValidationError
 from rest_framework import serializers
+from rest_framework.exceptions import ValidationError
 
 from care.facility.api.serializers.shifting import has_facility_permission
 from care.facility.models.facility import Facility
@@ -108,6 +109,7 @@ class FileUploadCreateSerializer(serializers.ModelSerializer):
     associating_id = serializers.CharField(write_only=True)
     internal_name = serializers.CharField(read_only=True)
     original_name = serializers.CharField(write_only=True)
+    mime_type = serializers.CharField(write_only=True)
 
     class Meta:
         model = FileUpload
@@ -120,19 +122,27 @@ class Meta:
             "signed_url",
             "internal_name",
             "original_name",
+            "mime_type",
         )
         write_only_fields = ("associating_id",)
 
     def create(self, validated_data):
         user = self.context["request"].user
+        mime_type = validated_data.pop("mime_type")
+
+        if mime_type not in settings.ALLOWED_MIME_TYPES:
+            raise ValidationError({"detail": "Invalid File Type"})
+
         internal_id = check_permissions(
             validated_data["file_type"], validated_data["associating_id"], user
         )
         validated_data["associating_id"] = internal_id
         validated_data["uploaded_by"] = user
         validated_data["internal_name"] = validated_data["original_name"]
         del validated_data["original_name"]
-        return super().create(validated_data)
+        file_upload = super().create(validated_data)
+        file_upload.signed_url = file_upload.signed_url(mime_type=mime_type)
+        return file_upload
 
 
 class FileUploadListSerializer(serializers.ModelSerializer):

care/facility/models/file_upload.py

@@ -83,14 +83,17 @@ def save(self, *args, **kwargs):
             self.internal_name = internal_name
         return super().save(*args, **kwargs)
 
-    def signed_url(self, duration=60 * 60):
+    def signed_url(self, duration=60 * 60, mime_type=None):
         s3Client = boto3.client("s3", **cs_provider.get_client_config())
+        params = {
+            "Bucket": settings.FILE_UPLOAD_BUCKET,
+            "Key": f"{self.FileType(self.file_type).name}/{self.internal_name}",
+        }
+        if mime_type:
+            params["ContentType"] = mime_type
         return s3Client.generate_presigned_url(
             "put_object",
-            Params={
-                "Bucket": settings.FILE_UPLOAD_BUCKET,
-                "Key": f"{self.FileType(self.file_type).name}/{self.internal_name}",
-            },
+            Params=params,
             ExpiresIn=duration,  # seconds
         )
 

care/facility/tests/test_file_upload.py

@@ -0,0 +1,48 @@
+from rest_framework import status
+from rest_framework.test import APITestCase
+
+from care.utils.tests.test_utils import TestUtils
+
+
+class FileUploadApiTestCase(TestUtils, APITestCase):
+    @classmethod
+    def setUpTestData(cls) -> None:
+        cls.state = cls.create_state()
+        cls.district = cls.create_district(cls.state)
+        cls.local_body = cls.create_local_body(cls.district)
+        cls.super_user = cls.create_super_user("su", cls.district)
+        cls.facility = cls.create_facility(cls.super_user, cls.district, cls.local_body)
+        cls.user = cls.create_user("nurse", cls.district, home_facility=cls.facility)
+        cls.patient = cls.create_patient(
+            cls.district, cls.facility, local_body=cls.local_body
+        )
+        cls.consultation = cls.create_consultation(cls.patient, cls.facility)
+
+    def test_file_upload_whitelist(self):
+        for mime_type in ["application/pdf", "image/png", "image/jpeg"]:
+            response = self.client.post(
+                "/api/v1/files/",
+                {
+                    "original_name": f"test.{mime_type.split('/')[-1]}",
+                    "file_type": "CONSULTATION",
+                    "name": "Test File",
+                    "associating_id": self.consultation.external_id,
+                    "file_category": "UNSPECIFIED",
+                    "mime_type": mime_type,
+                },
+            )
+            self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+        for mime_type in ["image/gif"]:
+            response = self.client.post(
+                "/api/v1/files/",
+                {
+                    "original_name": f"test.{mime_type.split('/')[-1]}",
+                    "file_type": "CONSULTATION",
+                    "name": "Test File",
+                    "associating_id": self.consultation.external_id,
+                    "file_category": "UNSPECIFIED",
+                    mime_type: mime_type,
+                },
+            )
+            self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)

config/settings/base.py

@@ -509,6 +509,45 @@
     default=f"https://{FILE_UPLOAD_BUCKET}.s3.amazonaws.com",
 )
 
+ALLOWED_MIME_TYPES = env.list(
+    "ALLOWED_MIME_TYPES",
+    default=[
+        # Images
+        "image/jpeg",
+        "image/png",
+        "image/gif",
+        "image/bmp",
+        "image/webp",
+        "image/svg+xml",
+        # Videos
+        "video/mp4",
+        "video/mpeg",
+        "video/x-msvideo",
+        "video/quicktime",
+        "video/x-ms-wmv",
+        "video/x-flv",
+        "video/webm",
+        # Audio
+        "audio/mpeg",
+        "audio/wav",
+        "audio/aac",
+        "audio/ogg",
+        "audio/midi",
+        "audio/x-midi",
+        "audio/webm",
+        # Documents
+        "text/plain",
+        "text/csv",
+        "application/rtf",
+        "application/msword",
+        "application/vnd.oasis.opendocument.text",
+        "application/pdf",
+        "application/vnd.ms-excel",
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        "application/vnd.oasis.opendocument.spreadsheet",
+    ],
+)
+
 FACILITY_S3_BUCKET = env("FACILITY_S3_BUCKET", default="")
 FACILITY_S3_REGION = env("FACILITY_S3_REGION_CODE", default="ap-south-1")
 FACILITY_S3_KEY = env("FACILITY_S3_KEY", default="")