Add more headers and conversation-id to Defender for AI (#858)

microsoft · Jun 19, 2024 · cd03c6a · cd03c6a
1 parent 7fac325
commit cd03c6a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/app.py b/app.py
@@ -207,7 +207,8 @@ def prepare_model_args(request_body, request_headers):
     user_json = None
     if (MS_DEFENDER_ENABLED):
         authenticated_user_details = get_authenticated_user_details(request_headers)
-        user_json = get_msdefender_user_json(authenticated_user_details, request_headers)
+        conversation_id = request_body.get("conversation_id", None)        
+        user_json = get_msdefender_user_json(authenticated_user_details, request_headers, conversation_id)
 
     model_args = {
         "messages": messages,

diff --git a/backend/security/ms_defender_utils.py b/backend/security/ms_defender_utils.py
@@ -1,11 +1,14 @@
 import json
 
-def get_msdefender_user_json(authenticated_user_details, request_headers):
+def get_msdefender_user_json(authenticated_user_details, request_headers, conversation_id):
     auth_provider = authenticated_user_details.get('auth_provider')
-    source_ip = request_headers.get('X-Forwarded-For', request_headers.get('Remote-Addr', ''))
+    source_ip = request_headers.get('Remote-Addr', '')
+    header_names = ['User-Agent', 'X-Forwarded-For', 'Forwarded', 'X-Real-IP', 'True-Client-IP', 'CF-Connecting-IP']
     user_args = {
         "EndUserId": authenticated_user_details.get('user_principal_id'),
         "EndUserIdType": "EntraId" if auth_provider == "aad" else auth_provider,
         "SourceIp": source_ip.split(':')[0], #remove port
+        "SourceRequestHeaders": {header: request_headers[header] for header in header_names if header in request_headers},
+        "ConversationId": conversation_id,
     }
     return json.dumps(user_args)