Update controllers/ws_controller.py

controllers/ws_controller.py

@@ -99,6 +99,25 @@ async def get_messages(code: str,websocket:WebSocket,start,brand_name: Optional[
     # print(messages)
     return messages
 
+def is_file_encrypted(filename: str, data: str) -> bool:
+    try:
+        file_content = io.BytesIO(base64.urlsafe_b64decode(data))
+        if filename.lower().endswith('.pdf'):
+            return is_pdf_encrypted(file_content)
+        # Add checks for other file types as needed
+    except Exception as e:
+        print(f"Error checking encryption for {filename}: {str(e)}")
+        return True  # Assume encrypted if there's an error
+    return False
+
+def is_pdf_encrypted(file_content):
+    try:
+        pdf_reader = PyPDF2.PdfReader(file_content)
+        return pdf_reader.is_encrypted
+    except:
+        return True  # Assume encrypted if there's an error
+
+
 async def process_message(message:Message, websocket:WebSocket, chunk_size:int):
     logging.info("process_message")
     print(message)
@@ -218,10 +237,22 @@ def extract_attachments_from_mail(access_token: str, message_data: dict) -> List
                 if filename.endswith(".zip") or filename.endswith(".txt") or filename.endswith(".png") or filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".gif"):
                     continue
                 data = attachment_data.get("data", "")
-                raw_text=ut.extract_text_from_attachment(filename , data)
-                struct_data = ut.strcuture_document_data(raw_text)
-                if struct_data:
-                    structured_data.append(struct_data)
+                if is_file_encrypted(filename, data):
+                    print(f"Skipping encrypted file: {filename}")
+                    continue  # Skip this file if it's encrypted
+                
+                try:
+                    raw_text = ut.extract_text_from_attachment(filename, data)
+                    struct_data = ut.strcuture_document_data(raw_text)
+                    if struct_data:
+                        structured_data.append(struct_data)
+                except Exception as e:
+                    print(f"Error processing attachment {filename}: {str(e)}")
+                    continue  # Skip this attachment if there's an error
+                # raw_text=ut.extract_text_from_attachment(filename , data)
+                # struct_data = ut.strcuture_document_data(raw_text)
+                # if struct_data:
+                #     structured_data.append(struct_data)
                 
                 attachments.append(Attachment(attachment_len = len(attachment_data.get("data", "")),filename=filename, data=attachment_data.get("data", "")))
     return attachments,structured_data