wip

Author: dominic-r

authentik/root/storages.py

@@ -296,24 +296,23 @@ def get_valid_name(self, name: str) -> str:
         return super().get_valid_name(name)
 
     def _normalize_name(self, name: str) -> str:
-        """Create tenant-aware S3 key"""
-        try:
-            name = self.get_valid_name(name)
-            normalized = safe_join(self.location, self.get_tenant_path(name))
-            LOGGER.debug("Normalized S3 key", original=name, normalized=normalized)
-            return normalized
-        except ValueError as e:
-            LOGGER.error(
-                "Invalid S3 key path detected", name=name, tenant=self.tenant_prefix, error=str(e)
-            )
-            raise SuspiciousOperation(f"Invalid path: {name}") from e
+        """Normalize file path for S3 storage"""
+        if ".." in name or name.startswith("/"):
+            raise SuspiciousOperation(f"Suspicious path: {name}")
+        normalized = self.get_tenant_path(name)
+        LOGGER.debug(
+            "Normalized S3 key",
+            original=name,
+            normalized=normalized,
+        )
+        return normalized
 
     def _randomize_filename(self, filename: str) -> str:
-        """Generate random filename"""
-        stem = uuid.uuid4().hex
-        suffix = Path(filename).suffix.lower()
-        tenant_hash = uuid.uuid5(uuid.NAMESPACE_DNS, self.tenant_prefix).hex[:8]
-        randomized = f"{tenant_hash}_{stem}{suffix}"
+        """Generate a randomized filename while preserving extension"""
+        name, ext = os.path.splitext(filename)
+        tenant_hash = str(uuid.uuid5(uuid.NAMESPACE_DNS, self.tenant_prefix))[:8]
+        random_uuid = str(uuid.uuid4())
+        randomized = f"{tenant_hash}_{random_uuid}{ext.lower()}"
         LOGGER.debug(
             "Randomized filename",
             original=filename,
@@ -323,51 +322,57 @@ def _randomize_filename(self, filename: str) -> str:
         return randomized
 
     def _save(self, name: str, content) -> str:
-        """Save file with secure random name and handle old file cleanup"""
-        try:
-            name = self.get_valid_name(name)
-            randomized_name = self._randomize_filename(name)
-            normalized_name = self._normalize_name(randomized_name)
-
-            self._file_mapping[name] = normalized_name
-
-            LOGGER.info(
-                "Saving file to S3",
-                original_name=name,
-                randomized_name=randomized_name,
-                normalized_name=normalized_name,
-                tenant=self.tenant_prefix,
-            )
+        """Save file to S3 with tenant isolation and random filename"""
+        randomized_name = self._randomize_filename(name)
+        normalized_name = self._normalize_name(randomized_name)
+
+        LOGGER.info(
+            "Saving file to S3",
+            original_name=name,
+            randomized_name=randomized_name,
+            normalized_name=normalized_name,
+            tenant=self.tenant_prefix,
+        )
 
-            result = super()._save(randomized_name, content)
+        try:
+            # Upload file to S3
+            self.bucket.Object(normalized_name).upload_fileobj(content)
 
+            # Verify upload
             try:
                 self.bucket.Object(normalized_name).load()
-                LOGGER.debug(
-                    "File saved successfully to S3", key=normalized_name, tenant=self.tenant_prefix
-                )
             except ClientError as e:
+                error_code = e.response.get("Error", {}).get("Code", "Unknown")
+                error_message = e.response.get("Error", {}).get("Message", "Unknown error")
                 LOGGER.error(
                     "Failed to verify S3 upload",
                     key=normalized_name,
-                    error_code=e.response["Error"]["Code"],
-                    message=e.response["Error"]["Message"],
+                    error_code=error_code,
+                    message=error_message,
                     tenant=self.tenant_prefix,
                 )
-                self._delete_file(normalized_name)
+                # Clean up failed upload
+                try:
+                    self.bucket.Object(normalized_name).delete()
+                except Exception as cleanup_error:
+                    LOGGER.error(
+                        "Failed to clean up failed upload",
+                        key=normalized_name,
+                        error=str(cleanup_error),
+                        tenant=self.tenant_prefix,
+                    )
                 raise
 
-            return result
+            # Store mapping of original name to normalized name
+            self._file_mapping[name] = normalized_name
 
-        except ClientError as e:
-            LOGGER.error(
-                "S3 upload failed",
-                error_code=e.response["Error"]["Code"],
-                message=e.response["Error"]["Message"],
+            LOGGER.debug(
+                "File saved successfully to S3",
                 key=normalized_name,
                 tenant=self.tenant_prefix,
             )
-            raise
+            return normalized_name
+
         except Exception as e:
             LOGGER.error(
                 "Unexpected error saving file to S3",
@@ -377,6 +382,34 @@ def _save(self, name: str, content) -> str:
             )
             raise
 
+    def delete(self, name: str) -> None:
+        """Delete file from S3"""
+        try:
+            # Get normalized name from mapping or normalize original name
+            normalized_name = self._file_mapping.get(name, self._normalize_name(name))
+            self.bucket.Object(normalized_name).delete()
+            # Remove from mapping if exists
+            self._file_mapping.pop(name, None)
+            LOGGER.debug(
+                "File deleted from S3",
+                key=normalized_name,
+                tenant=self.tenant_prefix,
+            )
+        except ClientError as e:
+            if e.response.get("Error", {}).get("Code") != "404":
+                LOGGER.error(
+                    "Failed to delete file from S3",
+                    name=name,
+                    error=str(e),
+                    tenant=self.tenant_prefix,
+                )
+                raise
+            LOGGER.debug(
+                "File not found during delete",
+                name=name,
+                tenant=self.tenant_prefix,
+            )
+
     def url(self, name: str, **kwargs) -> str:
         """Generate URL without signing parameters when using custom domain"""
         try: