Add encryption and decryption for sensitive message data

Integrated mechanisms to encrypt sensitive words in messages and decrypt them when retrieving. This ensures better data protection while handling sensitive user information in message flows.

test.py

@@ -1,13 +1,53 @@
-import asyncio
-
-import numpy as np
-
-from trauma.api.message.ai.openai_request import convert_value_to_embeddings
-from trauma.core.config import settings
-
-
-async def main():
-    entities = await settings.DB_CLIENT
-
-if __name__ == '__main__':
-    asyncio.run(main())
+# import re
+#
+# def xor_cipher(text: str, key: str) -> str:
+#     key_bytes = key.encode('utf-8')
+#     text_bytes = text.encode('utf-8')
+#     key_len = len(key_bytes)
+#     encrypted_bytes = bytes([
+#         text_bytes[i] ^ key_bytes[i % key_len]
+#         for i in range(len(text_bytes))
+#     ])
+#     return encrypted_bytes.hex()
+#
+# def encrypt_sensitive_data(text: str, words_to_encrypt: list[str], secret_key: str) -> str:
+#     result = text
+#     for word in words_to_encrypt:
+#         if word in result:
+#             encrypted = xor_cipher(word, secret_key)
+#             result = result.replace(word, f"[{encrypted}]")
+#     return result
+#
+# def decrypt_sensitive_data(text: str, secret_key: str) -> str:
+#
+#     def decrypt_match(match):
+#         encrypted_hex = match.group(1)
+#         # Convert hex back to bytes
+#         encrypted_bytes = bytes.fromhex(encrypted_hex)
+#         # XOR with key to decrypt
+#         key_bytes = secret_key.encode('utf-8')
+#         decrypted_bytes = bytes([
+#             encrypted_bytes[i] ^ key_bytes[i % len(key_bytes)]
+#             for i in range(len(encrypted_bytes))
+#         ])
+#         return decrypted_bytes.decode('utf-8')
+#
+#     # Find all [encrypted] patterns and decrypt them
+#     pattern = r'\[([\da-fA-F]+)\]'
+#     return re.sub(pattern, decrypt_match, text)
+#
+# # Example usage:
+# if __name__ == "__main__":
+#     SECRET_KEY = "dda7db64674d3cbc571ccedfdb4321818ba642b8dd3ddbdd80d1ce2b2a4a3546"
+#
+#     # Test encryption
+#     original_text = "Привет! Меня зовут John, я живу в Moscow, мой email: [email protected]"
+#     sensitive_words = []
+#
+#     encrypted_text = encrypt_sensitive_data(original_text, sensitive_words, SECRET_KEY)
+#     print("Encrypted:", encrypted_text)
+#
+#     # Test decryption
+#     decrypted_text = decrypt_sensitive_data(encrypted_text, SECRET_KEY)
+#     print("Decrypted:", decrypted_text)
+#

trauma/api/message/ai/engine.py

@@ -5,7 +5,7 @@ import numpy as np
 from trauma.api.chat.dto import EntityData
 from trauma.api.chat.model import ChatModel
 from trauma.api.data.model import EntityModel, EntityModelExtended
-from trauma.api.message.ai.openai_request import (update_entity_data_with_ai,
+from trauma.api.message.ai.openai_request import (get_sensitive_words, update_entity_data_with_ai,
                                                   generate_next_question,
                                                   generate_search_request,
                                                   generate_final_response,
@@ -28,25 +28,26 @@ from trauma.api.message.utils import (decode_treatment_letters,
                                       prepare_final_entities_str,
                                       pick_empty_field_instructions,
                                       find_matching_age_group,
-                                      search_changed_field_inst)
+                                      search_changed_field_inst,
+                                      encrypt_message)
 from trauma.core.config import settings
 
 
 async def search_entities(
         user_message: str, messages: list[MessageModel], chat: ChatModel
 ) -> CreateMessageResponse:
-    decoded_message = decode_treatment_letters(user_message)
-    message_history_str = prepare_message_history_str(messages, decoded_message)
+    user_message = decode_treatment_letters(user_message)
+    message_history_str = prepare_message_history_str(messages, user_message)
 
     entity_data, is_valid = await asyncio.gather(
-        update_entity_data_with_ai(chat.entityData, decoded_message, messages[-1].text),
-        check_is_valid_request(decoded_message, message_history_str)
+        update_entity_data_with_ai(chat.entityData, user_message, messages[-1].text),
+        check_is_valid_request(user_message, message_history_str)
     )
     final_entities, fields_changed_inst = None, search_changed_field_inst(entity_data, chat.entityData)
     
     if not is_valid:
         empty_field = retrieve_empty_field_from_entity_data(chat.entityData.model_dump(mode='json'))
-        response = await generate_invalid_response(decoded_message, message_history_str, empty_field)
+        response = await generate_invalid_response(user_message, message_history_str, empty_field)
         final_entities = messages[-1].entities if messages else None
 
     else:
@@ -57,7 +58,7 @@ async def search_entities(
         if empty_field == 'age':
             response = await generate_next_question(empty_field_instructions, message_history_str)
         else:
-            user_messages_str = prepare_user_messages_str(decoded_message, messages)
+            user_messages_str = prepare_user_messages_str(user_message, messages)
             possible_entity_indexes, search_request = await asyncio.gather(
                 filter_entities_by_age_location(entity_data),
                 generate_search_request(user_messages_str, entity_data)
@@ -68,16 +69,17 @@ async def search_entities(
             final_entities_str = prepare_final_entities_str(final_entities)
             if final_entities:
                 response = await generate_final_response(
-                    final_entities_str, decoded_message, message_history_str, empty_field_instructions
+                    final_entities_str, user_message, message_history_str, empty_field_instructions
                 )
             else:
                 response = await generate_empty_final_response(
-                    decoded_message, message_history_str, fields_changed_inst
+                    user_message, message_history_str, fields_changed_inst
                 )
 
-    user_message = MessageModel(chatId=chat.id, author=Author.User, text=decoded_message)
+    user_message = MessageModel(chatId=chat.id, author=Author.User, text=user_message)
     assistant_message = MessageModel(chatId=chat.id, author=Author.Assistant, text=response, entities=final_entities)
-    asyncio.create_task(save_assistant_user_message(user_message, assistant_message))
+    user_message_enc, assistant_message_enc = await encrypt_messages([user_message, assistant_message])
+    asyncio.create_task(save_assistant_user_message(user_message_enc, assistant_message_enc))
     return assistant_message
 
 
@@ -134,3 +136,13 @@ async def set_entities_score(entities: list[EntityModelExtended], search_request
         if score > 0.72:
             final_entities.append(entity)
     return sorted(final_entities, key=lambda x: x.score, reverse=True)
+
+
+async def encrypt_messages(messages: list[MessageModel]) -> list[MessageModel]:
+    encrypted_messages = []
+    sensitive_words = await asyncio.gather(*[get_sensitive_words(message.text) for message in messages])
+    for message, sensitive_word in zip(messages, sensitive_words):
+        encrypted_message = MessageModel(**message.model_dump())
+        encrypted_message.text = encrypt_message(message.text, sensitive_word)
+        encrypted_messages.append(encrypted_message)
+    return encrypted_messages

trauma/api/message/ai/openai_request.py

@@ -209,5 +209,13 @@ async def generate_searched_entity_response(user_query: str, facility: EntityMod
     return messages
 
 
-if __name__ == '__main__':
-    asyncio.run(retrieve_semantic_answer('I want to know more about Praktijk Hermens'))
+@openai_wrapper(is_json=True, return_='words')
+async def get_sensitive_words(text: str):
+    messages = [
+        {
+            "role": "system",
+            "content": TraumaPrompts.get_sensitive_words
+            .replace("{text}", text)
+        }
+    ]
+    return messages

trauma/api/message/ai/prompts.py

@@ -382,3 +382,33 @@ Your response must be in the following JSON format:
 }
 ```
 - **score**: A floating-point number between **0.00 and 1.00**, representing the degree of relevance."""
+    get_sensitive_words = """## Task
+
+You must find all sensitive words or word combinations in the text `{text}` and return them in a JSON object. Sensitive data includes:
+- Personal data
+- Name and surname
+- Email
+- Phone number
+- Date of birth
+- Address (street, house number, postal code)
+- IP-address
+- Civil registration number
+
+## Data
+
+**Text**:
+```
+{text}
+```
+
+## JSON Response Format
+
+```json
+{
+    "words": ["string", "string"]
+}
+```
+
+## Important Notes
+
+- The words must be in the same case and language as they appear in the text."""

trauma/api/message/utils.py

@@ -6,7 +6,7 @@ from trauma.api.data.dto import AgeGroup
 from trauma.api.data.model import EntityModel
 from trauma.api.message.dto import Author
 from trauma.api.message.model import MessageModel
-
+from trauma.core.config import settings
 
 def transform_messages_to_openai(messages: list[MessageModel]) -> list[dict]:
     openai_messages = []
@@ -268,3 +268,40 @@ def search_changed_field_inst(entity_data: dict, old_entity_data: EntityData) ->
             real_key = key if key!="treatmentArea" else "traumaType"
             changed_fields[real_key] = instruction_map[key]
     return changed_fields
+
+
+def xor_cipher(text: str, key: str) -> str:
+    key_bytes = key.encode('utf-8')
+    text_bytes = text.encode('utf-8')
+    key_len = len(key_bytes)
+    encrypted_bytes = bytes([
+        text_bytes[i] ^ key_bytes[i % key_len]
+        for i in range(len(text_bytes))
+    ])
+    return encrypted_bytes.hex()
+
+def encrypt_message(text: str, words_to_encrypt: list[str]) -> str:
+    result = text
+    for word in words_to_encrypt:
+        if word in result:
+            encrypted = xor_cipher(word, settings.SECRET_KEY)
+            result = result.replace(word, f"[{encrypted}]")
+    return result
+
+def decrypt_messages(messages: list[MessageModel]) -> list[MessageModel]:
+    
+    def decrypt_match(match):
+        encrypted_hex = match.group(1)
+        encrypted_bytes = bytes.fromhex(encrypted_hex)
+        key_bytes = settings.SECRET_KEY.encode('utf-8')
+        decrypted_bytes = bytes([
+            encrypted_bytes[i] ^ key_bytes[i % len(key_bytes)]
+            for i in range(len(encrypted_bytes))
+        ])
+        return decrypted_bytes.decode('utf-8')
+    
+    pattern = r'\[([\da-fA-F]+)\]'
+
+    for message in messages:
+        message.text = re.sub(pattern, decrypt_match, message.text)
+    return messages

trauma/api/message/views.py

@@ -13,7 +13,7 @@ from trauma.api.message.schemas import (AllMessageWrapper,
                                         CreateMessageRequest)
 from trauma.core.security import PermissionDependency
 from trauma.core.wrappers import TraumaResponseWrapper
-
+from trauma.api.message.utils import decrypt_messages
 
 @message_router.get('/{chatId}/all')
 async def get_all_chat_messages(
@@ -35,6 +35,7 @@ async def create_message(
         account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
 ) -> TraumaResponseWrapper[MessageModel]:
     messages, chat = await get_all_chat_messages_obj(chatId, account)
+    messages = decrypt_messages(messages)
     response = await search_entities(message_data.text, messages, chat)
     return TraumaResponseWrapper(data=response)