Raksith/ForensicRag
0
1import os
2import json
3import hashlib
4import logging
5import time
6from datetime import datetime, timedelta
7from typing import List, Dict, Any, Optional
8import threading
9import re
10
11from fastapi import FastAPI, HTTPException
12from fastapi.middleware.cors import CORSMiddleware
13from pydantic import BaseModel
14import uvicorn
15
16from pymongo import MongoClient, ASCENDING
17import numpy as np
18from sentence_transformers import SentenceTransformer
19from sklearn.metrics.pairwise import cosine_similarity
20import faiss
21
22import google.generativeai as genai
23import dotenv
24
25dotenv.load_dotenv()
26
27# Configure logging
28logging.basicConfig(level=logging.INFO)
29logger = logging.getLogger(__name__)
30
31# Environment variables
32MONGODB_URL = os.getenv("MONGODB_URL")
33MONGODB_NAME = os.getenv("MONGODB_NAME")
34MONGODB_COLLECTION = os.getenv("MONGODB_COLLECTION")
35GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
36
37# File paths
38DATA_DIR = os.getenv("DATA_DIR", "./data")
39os.makedirs(DATA_DIR, exist_ok=True)
40
41FAISS_INDEX_PATH = os.path.join(DATA_DIR, "faiss.index")
42DOCS_PATH = os.path.join(DATA_DIR, "docs.jsonl")
43EMBEDDINGS_PATH = os.path.join(DATA_DIR, "embeddings.json")
44
45# Initialize FastAPI
46app = FastAPI(title="UFDR Forensics RAG API", version="2.0.0")
47
48app.add_middleware(
49 CORSMiddleware,
50 allow_origins=["*"],
51 allow_credentials=True,
52 allow_methods=["*"],
53 allow_headers=["*"],
54)
55
56# Pydantic models
57class QueryRequest(BaseModel):
58 query: str
59 case_id: Optional[str] = None
60
61class QueryResponse(BaseModel):
62 query: str
63 direct_answer: str
64 case_info: Optional[Dict[str, Any]] = None
65 evidence_count: int
66 detected_case_id: Optional[str] = None
67
68class UFDRProcessor:
69 def __init__(self):
70 """Initialize UFDR processor with local storage and MongoDB"""
71 self.client = None
72 self.collection = None
73 self.local_embeddings = {}
74 self.local_documents = {}
75 self.doc_id_to_embeddings_map = {} # Maps doc_id to embedding indices
76
77 # Load local data
78 self._load_local_data()
79
80 # Initialize MongoDB connection
81 self._init_mongodb()
82
83 # Initialize embedding model
84 self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
85
86 # Initialize Gemini
87 self._init_gemini()
88
89 # Generate embeddings for any documents that don't have them
90 self._generate_missing_embeddings()
91
92 # Initialize FAISS index
93 self.faiss_index = None
94 self._load_or_create_index()
95
96 # Start monitoring thread
97 self.monitoring_active = True
98 self.monitoring_thread = threading.Thread(target=self._monitor_changes, daemon=True)
99 self.monitoring_thread.start()
100
101 logger.info("UFDR Processor initialized successfully")
102
103 def _init_mongodb(self):
104 """Initialize MongoDB connection"""
105 try:
106 if MONGODB_URL:
107 import certifi
108 ca = certifi.where()
109
110 self.client = MongoClient(
111 MONGODB_URL,
112 tls=True,
113 tlsCAFile=ca,
114 serverSelectionTimeoutMS=5000
115 )
116
117 self.client.admin.command("ping")
118 self.db = self.client[MONGODB_NAME]
119 self.collection = self.db[MONGODB_COLLECTION]
120
121 # Create indexes
122 self.collection.create_index([("$**", "text")])
123 self.collection.create_index([("reports.xml.UFDRReport.Case.@id", ASCENDING)])
124
125 logger.info("MongoDB connection successful")
126 else:
127 logger.warning("MongoDB URL not provided - running in local-only mode")
128 except Exception as e:
129 logger.error(f"MongoDB connection failed: {e}")
130 self.client = None
131
132 def _init_gemini(self):
133 """Initialize Gemini model"""
134 try:
135 if GEMINI_API_KEY:
136 genai.configure(api_key=GEMINI_API_KEY)
137 self.gemini_model = genai.GenerativeModel("gemini-1.5-pro")
138 logger.info("Gemini model initialized")
139 else:
140 self.gemini_model = None
141 logger.warning("Gemini API key not provided")
142 except Exception as e:
143 logger.error(f"Gemini initialization failed: {e}")
144 self.gemini_model = None
145
146 def _load_local_data(self):
147 """Load local embeddings and documents"""
148 try:
149 if os.path.exists(EMBEDDINGS_PATH):
150 with open(EMBEDDINGS_PATH, 'r') as f:
151 data = json.load(f)
152 self.local_embeddings = data.get('embeddings', {})
153 self.local_documents = data.get('documents', {})
154 logger.info(f"Loaded {len(self.local_documents)} documents from local storage")
155 else:
156 logger.info("No local data found, will sync from MongoDB")
157 except Exception as e:
158 logger.error(f"Error loading local data: {e}")
159 self.local_embeddings = {}
160 self.local_documents = {}
161
162 def _save_local_data(self):
163 """Save local data to file"""
164 try:
165 data = {
166 'embeddings': self.local_embeddings,
167 'documents': self.local_documents,
168 'last_updated': datetime.now().isoformat()
169 }
170 with open(EMBEDDINGS_PATH, 'w') as f:
171 json.dump(data, f, indent=2, default=str)
172 except Exception as e:
173 logger.error(f"Error saving local data: {e}")
174
175 def _process_ufdr_document(self, doc: Dict[str, Any]) -> Dict[str, Any]:
176 """Process UFDR document into searchable format with enhanced chunking"""
177 processed = {
178 'doc_id': str(doc.get('_id', '')),
179 'case_id': '',
180 'device': '',
181 'examiner': '',
182 'date': '',
183 'content_parts': []
184 }
185
186 # Extract case information
187 reports = doc.get('reports', {})
188 if 'xml' in reports and 'UFDRReport' in reports['xml']:
189 case_info = reports['xml']['UFDRReport'].get('Case', {})
190 processed['case_id'] = case_info.get('@id', '')
191 processed['device'] = case_info.get('@device', '')
192 processed['examiner'] = case_info.get('@examiner', '')
193 processed['date'] = case_info.get('@date', '')
194
195 if 'json' in reports:
196 json_data = reports['json']
197 if 'case' in json_data and not processed['case_id']:
198 processed['case_id'] = json_data['case']
199 if 'examiner' in json_data and not processed['examiner']:
200 processed['examiner'] = json_data['examiner']
201
202 # Extract artifacts with detailed metadata
203 artifacts = doc.get('artifacts', {})
204
205 # SMS messages - Enhanced with better context
206 sms_messages = artifacts.get('sms', [])
207 for msg in sms_messages:
208 sender = str(msg.get('sender', 'unknown'))
209 receiver = str(msg.get('receiver', 'unknown'))
210 message_text = msg.get('message', '')
211 timestamp = msg.get('timestamp', '')
212
213 content = f"SMS message from {sender} to {receiver} at {timestamp}: {message_text}"
214 processed['content_parts'].append({
215 'type': 'sms',
216 'content': content,
217 'timestamp': timestamp,
218 'metadata': msg,
219 'searchable_text': f"sms message text communication {sender} {receiver} {message_text}"
220 })
221
222 # Call logs - Enhanced
223 call_logs = artifacts.get('call_logs', [])
224 for call in call_logs:
225 caller = str(call.get('caller', 'unknown'))
226 callee = str(call.get('callee', 'unknown'))
227 duration = call.get('duration', 0)
228 timestamp = call.get('timestamp', '')
229
230 content = f"Call log: {caller} called {callee} for {duration} seconds at {timestamp}"
231 processed['content_parts'].append({
232 'type': 'call_log',
233 'content': content,
234 'timestamp': timestamp,
235 'metadata': call,
236 'searchable_text': f"call log phone call {caller} {callee} duration {duration}"
237 })
238
239 # Contacts - Enhanced
240 contacts = artifacts.get('contacts', [])
241 for contact in contacts:
242 name = contact.get('name', 'unknown')
243 phone = str(contact.get('phone', 'no phone'))
244
245 content = f"Contact saved: {name} with phone number {phone}"
246 processed['content_parts'].append({
247 'type': 'contact',
248 'content': content,
249 'metadata': contact,
250 'searchable_text': f"contact name {name} phone number {phone}"
251 })
252
253 # Locations - Enhanced
254 locations = artifacts.get('locations', [])
255 for location in locations:
256 lat = location.get('latitude', 'unknown')
257 lon = location.get('longitude', 'unknown')
258 timestamp = location.get('timestamp', '')
259
260 content = f"Location data at {timestamp}: coordinates {lat}, {lon}"
261 processed['content_parts'].append({
262 'type': 'location',
263 'content': content,
264 'timestamp': timestamp,
265 'metadata': location,
266 'searchable_text': f"location gps coordinates tracking {lat} {lon}"
267 })
268
269 # WiFi logs
270 wifi_logs = artifacts.get('wifi_logs', [])
271 for wifi in wifi_logs:
272 ssid = wifi.get('ssid', 'unknown')
273 mac = wifi.get('mac', 'unknown')
274 timestamp = wifi.get('timestamp', '')
275
276 content = f"WiFi connection to {ssid} (MAC: {mac}) at {timestamp}"
277 processed['content_parts'].append({
278 'type': 'wifi_log',
279 'content': content,
280 'timestamp': timestamp,
281 'metadata': wifi,
282 'searchable_text': f"wifi network connection {ssid} {mac}"
283 })
284
285 # Media files
286 media = doc.get('media', [])
287 for media_file in media:
288 path = media_file.get('path', 'unknown')
289 size = media_file.get('size_bytes', 0)
290 md5 = media_file.get('md5', '')
291
292 media_type = 'unknown'
293 if 'image' in path.lower() or any(ext in path.lower() for ext in ['.jpg', '.png', '.jpeg']):
294 media_type = 'image'
295 elif 'audio' in path.lower() or any(ext in path.lower() for ext in ['.mp3', '.amr', '.wav']):
296 media_type = 'audio'
297 elif 'video' in path.lower() or any(ext in path.lower() for ext in ['.mp4', '.avi']):
298 media_type = 'video'
299
300 content = f"Media file ({media_type}): {path}, size {size} bytes, MD5: {md5}"
301 processed['content_parts'].append({
302 'type': 'media',
303 'content': content,
304 'metadata': media_file,
305 'searchable_text': f"media file {media_type} {path} photo image audio video"
306 })
307
308 # Apps data
309 apps = doc.get('apps', {})
310 for app_name, app_data in apps.items():
311 content = f"Application data recovered: {app_name}"
312 processed['content_parts'].append({
313 'type': 'app_data',
314 'content': content,
315 'metadata': {'app_name': app_name, 'data': app_data},
316 'searchable_text': f"app application {app_name} data"
317 })
318
319 # Errors and anomalies
320 errors = doc.get('errors', [])
321 if errors:
322 content = f"Anomalies detected: {', '.join(str(e) for e in errors)}"
323 processed['content_parts'].append({
324 'type': 'anomaly',
325 'content': content,
326 'metadata': {'errors': errors},
327 'searchable_text': f"anomaly error issue problem suspicious {' '.join(str(e) for e in errors)}"
328 })
329
330 # Hash verification issues
331 hashes = doc.get('hashes', {})
332 if 'verification' in hashes:
333 for file_name, verification in hashes['verification'].items():
334 if not verification.get('exists', True):
335 content = f"Anomaly: File {file_name} declared but not found"
336 processed['content_parts'].append({
337 'type': 'anomaly',
338 'content': content,
339 'metadata': verification,
340 'searchable_text': f"anomaly missing file hash verification {file_name}"
341 })
342
343 return processed
344
345 def _generate_embeddings_for_document(self, doc_id: str, doc_data: Dict[str, Any]):
346 """Generate embeddings for document content with enhanced semantic chunks"""
347 try:
348 embeddings_list = []
349
350 case_id = doc_data.get('case_id', 'unknown')
351 device = doc_data.get('device', 'unknown')
352 examiner = doc_data.get('examiner', 'unknown')
353 date = doc_data.get('date', 'unknown')
354
355 # Main document embedding with multiple variants for better matching
356 main_contents = [
357 f"Case {case_id} forensic investigation device {device} examined by {examiner} on {date}",
358 f"Forensic case {case_id} phone device {device}",
359 f"Digital forensics case {case_id} examiner {examiner}",
360 f"Case identifier {case_id} investigation details"
361 ]
362
363 for content in main_contents:
364 embedding = self.embedding_model.encode(content)
365 embeddings_list.append({
366 'content': content,
367 'type': 'case_info',
368 'embedding': embedding.tolist(),
369 'metadata': {
370 'case_id': case_id,
371 'device': device,
372 'examiner': examiner,
373 'date': date
374 }
375 })
376
377 # Summary embedding for evidence count queries
378 artifact_count = len(doc_data.get('content_parts', []))
379 summary_content = f"Case {case_id} contains {artifact_count} pieces of evidence including artifacts and forensic data"
380 embedding = self.embedding_model.encode(summary_content)
381 embeddings_list.append({
382 'content': summary_content,
383 'type': 'summary',
384 'embedding': embedding.tolist(),
385 'metadata': {
386 'case_id': case_id,
387 'artifact_count': artifact_count
388 }
389 })
390
391 # Content part embeddings with enhanced searchable text
392 for part in doc_data.get('content_parts', []):
393 content = part.get('content', '')
394 searchable = part.get('searchable_text', content)
395
396 if content:
397 # Primary embedding from natural content
398 embedding = self.embedding_model.encode(content)
399 embeddings_list.append({
400 'content': content,
401 'type': part.get('type', 'unknown'),
402 'embedding': embedding.tolist(),
403 'metadata': part.get('metadata', {}),
404 'timestamp': part.get('timestamp', '')
405 })
406
407 # Secondary embedding from searchable text for better keyword matching
408 if searchable != content:
409 embedding2 = self.embedding_model.encode(searchable)
410 embeddings_list.append({
411 'content': content, # Keep original content
412 'type': part.get('type', 'unknown'),
413 'embedding': embedding2.tolist(),
414 'metadata': part.get('metadata', {}),
415 'timestamp': part.get('timestamp', ''),
416 'is_searchable_variant': True
417 })
418
419 self.local_embeddings[doc_id] = embeddings_list
420 logger.info(f"Generated {len(embeddings_list)} embeddings for document {doc_id}")
421
422 except Exception as e:
423 logger.error(f"Error generating embeddings for document {doc_id}: {e}")
424
425 def _generate_missing_embeddings(self):
426 """Generate embeddings for documents that don't have them"""
427 try:
428 for doc_id, doc_data in self.local_documents.items():
429 if doc_id not in self.local_embeddings:
430 logger.info(f"Generating embeddings for document {doc_id}")
431 self._generate_embeddings_for_document(doc_id, doc_data)
432
433 if self.local_embeddings:
434 self._save_local_data()
435 logger.info("Updated embeddings saved")
436
437 except Exception as e:
438 logger.error(f"Error generating missing embeddings: {e}")
439
440 def _load_or_create_index(self):
441 """Load existing FAISS index or create new one"""
442 try:
443 if os.path.exists(FAISS_INDEX_PATH) and self.local_embeddings:
444 self.faiss_index = faiss.read_index(FAISS_INDEX_PATH)
445 logger.info("Loaded existing FAISS index")
446 else:
447 self._rebuild_index()
448 except Exception as e:
449 logger.error(f"Error loading FAISS index: {e}")
450 self._rebuild_index()
451
452 def _rebuild_index(self):
453 """Rebuild FAISS index from documents"""
454 try:
455 if self.collection is not None:
456 self._sync_from_mongodb()
457
458 if not self.local_documents:
459 logger.warning("No documents available for indexing")
460 return
461
462 embeddings_list = []
463 doc_ids = []
464
465 for doc_id, doc_data in self.local_documents.items():
466 if doc_id in self.local_embeddings:
467 for embedding_data in self.local_embeddings[doc_id]:
468 embeddings_list.append(embedding_data['embedding'])
469 doc_ids.append(doc_id)
470 else:
471 self._generate_embeddings_for_document(doc_id, doc_data)
472 if doc_id in self.local_embeddings:
473 for embedding_data in self.local_embeddings[doc_id]:
474 embeddings_list.append(embedding_data['embedding'])
475 doc_ids.append(doc_id)
476
477 if embeddings_list:
478 embeddings_array = np.array(embeddings_list, dtype=np.float32)
479 dimension = embeddings_array.shape[1]
480
481 self.faiss_index = faiss.IndexFlatIP(dimension)
482 self.faiss_index.add(embeddings_array)
483
484 faiss.write_index(self.faiss_index, FAISS_INDEX_PATH)
485
486 logger.info(f"Rebuilt FAISS index with {len(embeddings_list)} embeddings")
487
488 except Exception as e:
489 logger.error(f"Error rebuilding index: {e}")
490
491 def _sync_from_mongodb(self):
492 """Sync documents from MongoDB to local storage"""
493 try:
494 if self.collection is None:
495 return
496
497 documents = list(self.collection.find({}))
498 for doc in documents:
499 doc_id = str(doc['_id'])
500 processed_doc = self._process_ufdr_document(doc)
501 self.local_documents[doc_id] = processed_doc
502
503 if doc_id not in self.local_embeddings:
504 self._generate_embeddings_for_document(doc_id, processed_doc)
505
506 self._save_local_data()
507 logger.info(f"Synced {len(documents)} documents from MongoDB")
508
509 except Exception as e:
510 logger.error(f"Error syncing from MongoDB: {e}")
511
512 def _update_single_document(self, doc: Dict[str, Any]):
513 """Update a single document in local storage and index efficiently"""
514 try:
515 doc_id = str(doc['_id'])
516 processed_doc = self._process_ufdr_document(doc)
517 self.local_documents[doc_id] = processed_doc
518 self._generate_embeddings_for_document(doc_id, processed_doc)
519 self._save_local_data()
520 self._rebuild_index()
521 logger.info(f"Successfully updated document {doc_id}")
522 except Exception as e:
523 logger.error(f"Error updating single document: {e}")
524
525 def _monitor_changes(self):
526 """Monitor MongoDB for changes"""
527 last_doc_count = 0
528 last_update_time = datetime.now()
529
530 if self.collection is not None:
531 try:
532 last_doc_count = self.collection.count_documents({})
533 logger.info(f"Initial document count: {last_doc_count}")
534 except Exception as e:
535 logger.error(f"Error getting initial count: {e}")
536
537 while self.monitoring_active:
538 try:
539 if self.collection is None:
540 time.sleep(5)
541 continue
542
543 try:
544 with self.collection.watch() as stream:
545 logger.info("Change stream established")
546
547 for change in stream:
548 if not self.monitoring_active:
549 break
550
551 operation = change.get('operationType', '')
552 doc_key = change.get('documentKey', {})
553 doc_id = str(doc_key.get('_id', 'unknown'))
554
555 logger.info(f"Change detected: {operation} on document {doc_id}")
556
557 if operation in ['insert', 'update', 'replace']:
558 doc = change.get('fullDocument') or self.collection.find_one({'_id': doc_key['_id']})
559 if doc:
560 self._update_single_document(doc)
561
562 elif operation == 'delete':
563 if doc_id in self.local_documents:
564 del self.local_documents[doc_id]
565 if doc_id in self.local_embeddings:
566 del self.local_embeddings[doc_id]
567 self._save_local_data()
568 self._rebuild_index()
569
570 time.sleep(0.5)
571
572 except Exception as stream_error:
573 logger.warning(f"Change stream failed: {stream_error}. Falling back to polling.")
574
575 while self.monitoring_active:
576 try:
577 current_count = self.collection.count_documents({})
578
579 if current_count != last_doc_count:
580 logger.info(f"Document count changed: {last_doc_count} -> {current_count}")
581 self._sync_from_mongodb()
582 self._rebuild_index()
583 last_doc_count = current_count
584 last_update_time = datetime.now()
585
586 time.sleep(2)
587
588 except Exception as poll_error:
589 logger.error(f"Polling error: {poll_error}")
590 time.sleep(5)
591
592 except Exception as e:
593 logger.error(f"Critical monitoring error: {e}")
594 time.sleep(10)
595
596 def _extract_case_id_from_query(self, query: str) -> Optional[str]:
597 """Extract case ID from query with multiple patterns"""
598 patterns = [
599 r'case\s*(?:id\s*)?[:\-]?\s*(\d+)', # "case 004", "case: 004", "case id 004"
600 r'for\s+case\s+(\d+)', # "for case 004"
601 r'in\s+case\s+(\d+)', # "in case 004"
602 r'from\s+case\s+(\d+)', # "from case 004"
603 r'of\s+case\s+(\d+)', # "of case 004"
604 r'\bcase\s+(\d+)\b', # "case 004" as whole word
605 r'for\s+(\d{3})\b', # "for 004"
606 r'\b(\d{3})\b', # standalone "004"
607 ]
608
609 for pattern in patterns:
610 match = re.search(pattern, query, re.IGNORECASE)
611 if match:
612 case_id = match.group(1).zfill(3) # Pad with zeros
613 logger.info(f"Extracted case ID: {case_id}")
614 return case_id
615
616 return None
617
618 def search_documents(self, query: str, case_id: str = None, limit: int = 15) -> List[Dict[str, Any]]:
619 """Enhanced semantic search with better relevance"""
620 try:
621 if not self.local_embeddings:
622 logger.warning("No embeddings available")
623 return []
624
625 query_embedding = self.embedding_model.encode(query)
626
627 all_embeddings = []
628 embedding_metadata = []
629
630 for doc_id, embeddings_list in self.local_embeddings.items():
631 doc_data = self.local_documents.get(doc_id, {})
632
633 if case_id:
634 doc_case_id = doc_data.get('case_id', '').strip()
635 if doc_case_id != case_id:
636 continue
637
638 for embedding_data in embeddings_list:
639 all_embeddings.append(embedding_data['embedding'])
640 embedding_metadata.append({
641 'doc_id': doc_id,
642 'content': embedding_data['content'],
643 'type': embedding_data['type'],
644 'metadata': embedding_data.get('metadata', {}),
645 'timestamp': embedding_data.get('timestamp', ''),
646 'case_info': {
647 'case_id': doc_data.get('case_id', ''),
648 'device': doc_data.get('device', ''),
649 'examiner': doc_data.get('examiner', ''),
650 'date': doc_data.get('date', '')
651 }
652 })
653
654 if not all_embeddings:
655 return []
656
657 embeddings_array = np.array(all_embeddings, dtype=np.float32)
658 query_embedding = query_embedding.astype(np.float32).reshape(1, -1)
659
660 similarities = cosine_similarity(query_embedding, embeddings_array)[0]
661
662 results = []
663 for i, similarity in enumerate(similarities):
664 if similarity > 0.05: # Lower threshold for more results
665 result = embedding_metadata[i].copy()
666 result['similarity'] = float(similarity)
667 results.append(result)
668
669 results.sort(key=lambda x: x['similarity'], reverse=True)
670 return results[:limit]
671
672 except Exception as e:
673 logger.error(f"Error in search: {e}")
674 return []
675
676 def generate_direct_answer(self, query: str, evidence: List[Dict[str, Any]], case_id: str = None) -> str:
677 """Generate enhanced direct answer with better context understanding"""
678 try:
679 if not evidence:
680 if case_id:
681 return f"No evidence found for case {case_id}."
682 return "No relevant evidence found in the forensic database."
683
684 query_lower = query.lower()
685
686 # Get all evidence for the case to provide comprehensive answers
687 case_evidence = [e for e in evidence if e['case_info'].get('case_id') == case_id] if case_id else evidence
688
689 # Count queries
690 if any(word in query_lower for word in ['how many', 'count', 'number of', 'total']):
691 if 'evidence' in query_lower or 'piece' in query_lower or 'artifact' in query_lower:
692 # Count unique artifacts by type
693 artifact_types = {}
694 for item in case_evidence:
695 art_type = item.get('type', 'unknown')
696 if art_type not in ['case_info', 'summary']:
697 artifact_types[art_type] = artifact_types.get(art_type, 0) + 1
698
699 total = sum(artifact_types.values())
700 if total > 0:
701 breakdown = ', '.join([f"{count} {type}" for type, count in artifact_types.items()])
702 return f"{total} pieces of evidence were recovered from case {case_id}: {breakdown}"
703
704 # Device queries
705 if 'device' in query_lower or 'phone' in query_lower:
706 for item in evidence:
707 if item.get('type') == 'case_info':
708 device = item['case_info'].get('device', '')
709 if device:
710 return device
711
712 # Examiner queries
713 if 'examiner' in query_lower or 'examined' in query_lower or 'who examined' in query_lower:
714 for item in evidence:
715 if item.get('type') == 'case_info':
716 examiner = item['case_info'].get('examiner', '')
717 date = item['case_info'].get('date', '')
718 if examiner:
719 if date:
720 return f"{examiner} examined the device on {date}"
721 return examiner
722
723 # SMS queries
724 if 'sms' in query_lower or 'message' in query_lower or 'text' in query_lower:
725 sms_items = [e for e in case_evidence if e.get('type') == 'sms']
726 if sms_items:
727 if len(sms_items) == 1:
728 return sms_items[0]['content']
729 else:
730 messages = [f"- {item['content']}" for item in sms_items[:5]]
731 return f"Found {len(sms_items)} SMS messages:\n" + "\n".join(messages)
732
733 # Call log queries
734 if 'call' in query_lower:
735 call_items = [e for e in case_evidence if e.get('type') == 'call_log']
736 if call_items:
737 if len(call_items) == 1:
738 return call_items[0]['content']
739 else:
740 calls = [f"- {item['content']}" for item in call_items[:5]]
741 return f"Found {len(call_items)} call logs:\n" + "\n".join(calls)
742
743 # Contact queries
744 if 'contact' in query_lower:
745 contact_items = [e for e in case_evidence if e.get('type') == 'contact']
746 if contact_items:
747 contacts = [f"- {item['content']}" for item in contact_items]
748 return f"Found {len(contact_items)} contacts:\n" + "\n".join(contacts)
749
750 # Location queries
751 if 'location' in query_lower or 'gps' in query_lower or 'tracking' in query_lower:
752 location_items = [e for e in case_evidence if e.get('type') == 'location']
753 if location_items:
754 locations = [f"- {item['content']}" for item in location_items]
755 return f"Found {len(location_items)} location points:\n" + "\n".join(locations)
756
757 # Media queries
758 if 'media' in query_lower or 'file' in query_lower or 'image' in query_lower or 'photo' in query_lower:
759 media_items = [e for e in case_evidence if e.get('type') == 'media']
760 if media_items:
761 media_list = [f"- {item['content']}" for item in media_items]
762 return f"Found {len(media_items)} media files:\n" + "\n".join(media_list)
763
764 # Anomaly/suspicious activity queries
765 if any(word in query_lower for word in ['anomaly', 'anomalies', 'suspicious', 'issue', 'error', 'problem']):
766 # First check if asking about a specific case
767 query_case_id = case_id
768 if not query_case_id:
769 query_case_id = self._extract_case_id_from_query(query)
770
771 if query_case_id:
772 # Check anomalies for specific case
773 case_evidence_filtered = [e for e in evidence if e['case_info'].get('case_id') == query_case_id]
774 anomaly_items = [e for e in case_evidence_filtered if e.get('type') == 'anomaly']
775
776 if anomaly_items:
777 anomalies = [f" - {item['content']}" for item in anomaly_items]
778 return f"Anomalies detected in case {query_case_id}:\n" + "\n".join(anomalies)
779 else:
780 return f"No anomalies detected in case {query_case_id}. All evidence integrity checks passed."
781 else:
782 # Global anomaly check across all cases
783 all_anomalies = {}
784 for doc_id, doc_data in self.local_documents.items():
785 doc_case_id = doc_data.get('case_id', 'unknown')
786 for part in doc_data.get('content_parts', []):
787 if part.get('type') == 'anomaly':
788 if doc_case_id not in all_anomalies:
789 all_anomalies[doc_case_id] = []
790 all_anomalies[doc_case_id].append(part['content'])
791
792 if all_anomalies:
793 result = "Anomalies detected across cases:\n\n"
794 for case_id_key, anomalies in all_anomalies.items():
795 result += f"Case {case_id_key}:\n"
796 for anomaly in anomalies:
797 result += f" - {anomaly}\n"
798 result += "\n"
799 return result.strip()
800 else:
801 return "No anomalies detected in any cases. All forensic data appears intact."
802
803 # "All details" or comprehensive queries
804 if any(phrase in query_lower for phrase in ['all details', 'everything', 'full details', 'complete']):
805 case_info = evidence[0]['case_info']
806 summary = f"Case {case_info['case_id']} - Complete Forensic Report\n{'='*60}\n\n"
807 summary += f"CASE INFORMATION:\n"
808 summary += f" Case ID: {case_info['case_id']}\n"
809 summary += f" Device: {case_info['device']}\n"
810 summary += f" Examiner: {case_info['examiner']}\n"
811 summary += f" Date: {case_info['date']}\n\n"
812
813 # Group by type with full details
814 by_type = {}
815 for item in case_evidence:
816 item_type = item.get('type', 'unknown')
817 if item_type not in ['case_info', 'summary', 'is_searchable_variant']:
818 if item_type not in by_type:
819 by_type[item_type] = []
820 by_type[item_type].append(item)
821
822 # Display each category with full details
823 if 'sms' in by_type:
824 summary += f"SMS MESSAGES ({len(by_type['sms'])}):\n"
825 for i, item in enumerate(by_type['sms'], 1):
826 summary += f" {i}. {item['content']}\n"
827 summary += "\n"
828
829 if 'call_log' in by_type:
830 summary += f"CALL LOGS ({len(by_type['call_log'])}):\n"
831 for i, item in enumerate(by_type['call_log'], 1):
832 summary += f" {i}. {item['content']}\n"
833 summary += "\n"
834
835 if 'contact' in by_type:
836 summary += f"CONTACTS ({len(by_type['contact'])}):\n"
837 for i, item in enumerate(by_type['contact'], 1):
838 summary += f" {i}. {item['content']}\n"
839 summary += "\n"
840
841 if 'location' in by_type:
842 summary += f"LOCATION DATA ({len(by_type['location'])}):\n"
843 for i, item in enumerate(by_type['location'], 1):
844 summary += f" {i}. {item['content']}\n"
845 summary += "\n"
846
847 if 'wifi_log' in by_type:
848 summary += f"WIFI CONNECTIONS ({len(by_type['wifi_log'])}):\n"
849 for i, item in enumerate(by_type['wifi_log'], 1):
850 summary += f" {i}. {item['content']}\n"
851 summary += "\n"
852
853 if 'media' in by_type:
854 summary += f"MEDIA FILES ({len(by_type['media'])}):\n"
855 for i, item in enumerate(by_type['media'], 1):
856 summary += f" {i}. {item['content']}\n"
857 summary += "\n"
858
859 if 'app_data' in by_type:
860 summary += f"APPLICATION DATA ({len(by_type['app_data'])}):\n"
861 for i, item in enumerate(by_type['app_data'], 1):
862 summary += f" {i}. {item['content']}\n"
863 summary += "\n"
864
865 if 'anomaly' in by_type:
866 summary += f"ANOMALIES DETECTED ({len(by_type['anomaly'])}):\n"
867 for i, item in enumerate(by_type['anomaly'], 1):
868 summary += f" {i}. {item['content']}\n"
869 summary += "\n"
870 else:
871 summary += "ANOMALIES: None detected\n\n"
872
873 if 'hash_declared' in by_type:
874 summary += f"HASH VERIFICATION ({len(by_type['hash_declared'])}):\n"
875 for i, item in enumerate(by_type['hash_declared'], 1):
876 summary += f" {i}. {item['content']}\n"
877 summary += "\n"
878
879 # Add total count
880 total_artifacts = sum(len(items) for items in by_type.values())
881 summary += f"{'='*60}\n"
882 summary += f"TOTAL EVIDENCE ITEMS: {total_artifacts}\n"
883
884 return summary
885
886 # Use Gemini for complex queries
887 if self.gemini_model:
888 context = []
889 for item in evidence[:5]:
890 context.append(f"{item['type']}: {item['content']}")
891
892 context_text = "\n".join(context)
893
894 prompt = f"""You are a digital forensics expert analyzing UFDR evidence. Provide a direct, factual answer.
895
896 Query: {query}
897 Case ID: {case_id or 'Not specified'}
898
899 Evidence:
900 {context_text}
901
902 Instructions:
903 1. Answer directly and concisely based ONLY on the evidence provided
904 2. For "show me" queries, list all relevant items
905 3. For count queries, provide exact numbers with breakdown
906 4. For specific data queries (SMS, calls, etc.), show the actual data
907 5. Keep formatting clean and readable
908 6. Maximum 200 words
909 7. Do NOT add explanations beyond what's in the evidence
910
911 Answer:"""
912
913 try:
914 response = self.gemini_model.generate_content(
915 prompt,
916 generation_config=genai.types.GenerationConfig(
917 temperature=0.2,
918 max_output_tokens=300,
919 )
920 )
921
922 if response and hasattr(response, 'text') and response.text:
923 return response.text.strip()
924
925 except Exception as e:
926 logger.warning(f"Gemini API error: {e}")
927
928 # Fallback: return most relevant evidence
929 if evidence:
930 top_items = evidence[:3]
931 result_parts = []
932 for item in top_items:
933 result_parts.append(item['content'])
934 return "\n".join(result_parts)
935
936 return "Unable to generate answer from available evidence."
937
938 except Exception as e:
939 logger.error(f"Error generating answer: {e}")
940 return "Error processing forensic evidence."
941
942 def query_system(self, query: str, case_id: str = None) -> Dict[str, Any]:
943 """Main query interface with enhanced processing"""
944 try:
945 # Extract case ID from query if not provided
946 detected_case_id = None
947 if not case_id:
948 detected_case_id = self._extract_case_id_from_query(query)
949 if detected_case_id:
950 case_id = detected_case_id
951 logger.info(f"Auto-extracted case_id: {case_id}")
952
953 logger.info(f"Processing query: '{query}' with case_id: {case_id}")
954
955 # Search for evidence
956 evidence = self.search_documents(query, case_id, limit=20)
957 logger.info(f"Found {len(evidence)} evidence items")
958
959 # Log top results for debugging
960 if evidence:
961 logger.info(f"Top result: type={evidence[0].get('type')}, similarity={evidence[0].get('similarity'):.3f}")
962
963 # Generate direct answer
964 direct_answer = self.generate_direct_answer(query, evidence, case_id)
965 logger.info(f"Generated answer length: {len(direct_answer)} chars")
966
967 # Extract case information
968 case_info = None
969 if evidence:
970 case_info = evidence[0]['case_info']
971
972 return {
973 'query': query,
974 'direct_answer': direct_answer,
975 'case_info': case_info,
976 'evidence_count': len(evidence),
977 'evidence': evidence[:5], # Return top 5 for debugging
978 'detected_case_id': detected_case_id
979 }
980
981 except Exception as e:
982 logger.error(f"Error in query system: {e}", exc_info=True)
983 return {
984 'query': query,
985 'direct_answer': f'Error processing query: {str(e)}',
986 'case_info': None,
987 'evidence_count': 0,
988 'evidence': [],
989 'detected_case_id': None
990 }
991
992# Initialize processor
993processor = UFDRProcessor()
994
995@app.get("/")
996async def read_root():
997 """Provides a welcome message and basic API information."""
998 return {
999 "message": "Welcome to the UFDR Forensics RAG API",
1000 "status": "online",
1001 "documentation": "/docs"
1002 }
1003
1004# API Endpoints
1005@app.post("/query", response_model=QueryResponse)
1006async def query_forensics(request: QueryRequest):
1007 """Main query endpoint for UFDR forensics analysis"""
1008 try:
1009 result = processor.query_system(request.query, request.case_id)
1010
1011 return QueryResponse(
1012 query=result['query'],
1013 direct_answer=result['direct_answer'],
1014 case_info=result['case_info'],
1015 evidence_count=result['evidence_count'],
1016 detected_case_id=result.get('detected_case_id')
1017 )
1018
1019 except Exception as e:
1020 logger.error(f"Error processing query: {e}")
1021 raise HTTPException(status_code=500, detail=str(e))
1022
1023@app.get("/status")
1024async def get_status():
1025 """Get comprehensive system status"""
1026 try:
1027 mongodb_status = "disconnected"
1028 mongodb_doc_count = 0
1029
1030 if processor.client:
1031 try:
1032 processor.client.admin.command("ping")
1033 mongodb_status = "connected"
1034 if processor.collection:
1035 mongodb_doc_count = processor.collection.count_documents({})
1036 except Exception:
1037 mongodb_status = "connection_error"
1038
1039 return {
1040 "status": "healthy",
1041 "timestamp": datetime.now().isoformat(),
1042 "mongodb": {
1043 "status": mongodb_status,
1044 "documents_in_db": mongodb_doc_count
1045 },
1046 "local_storage": {
1047 "documents_loaded": len(processor.local_documents),
1048 "embeddings_generated": sum(len(embs) for embs in processor.local_embeddings.values())
1049 },
1050 "search_engine": {
1051 "faiss_index_ready": processor.faiss_index is not None,
1052 "embedding_model": "all-MiniLM-L6-v2"
1053 },
1054 "ai_services": {
1055 "gemini_available": processor.gemini_model is not None
1056 },
1057 "monitoring": {
1058 "auto_sync_active": processor.monitoring_active,
1059 "real_time_updates": "enabled"
1060 }
1061 }
1062
1063 except Exception as e:
1064 return {
1065 "status": "error",
1066 "error": str(e),
1067 "timestamp": datetime.now().isoformat()
1068 }
1069
1070@app.get("/cases")
1071async def list_cases():
1072 """List all available forensic cases"""
1073 try:
1074 cases = []
1075 case_stats = {}
1076
1077 for doc_id, doc_data in processor.local_documents.items():
1078 case_id = doc_data.get('case_id', '')
1079 if case_id:
1080 case_info = {
1081 'case_id': case_id,
1082 'device': doc_data.get('device', ''),
1083 'examiner': doc_data.get('examiner', ''),
1084 'date': doc_data.get('date', ''),
1085 'artifacts_count': len(doc_data.get('content_parts', [])),
1086 'doc_id': doc_id
1087 }
1088
1089 artifact_types = {}
1090 for part in doc_data.get('content_parts', []):
1091 part_type = part.get('type', 'unknown')
1092 artifact_types[part_type] = artifact_types.get(part_type, 0) + 1
1093
1094 case_info['artifact_types'] = artifact_types
1095 cases.append(case_info)
1096
1097 if case_id not in case_stats:
1098 case_stats[case_id] = {
1099 'total_artifacts': 0,
1100 'devices': set(),
1101 'examiners': set()
1102 }
1103
1104 case_stats[case_id]['total_artifacts'] += case_info['artifacts_count']
1105 if case_info['device']:
1106 case_stats[case_id]['devices'].add(case_info['device'])
1107 if case_info['examiner']:
1108 case_stats[case_id]['examiners'].add(case_info['examiner'])
1109
1110 for case_id, stats in case_stats.items():
1111 stats['devices'] = list(stats['devices'])
1112 stats['examiners'] = list(stats['examiners'])
1113
1114 return {
1115 "cases": cases,
1116 "total_cases": len(set(case['case_id'] for case in cases)),
1117 "total_documents": len(cases),
1118 "case_statistics": case_stats
1119 }
1120
1121 except Exception as e:
1122 raise HTTPException(status_code=500, detail=f"Error listing cases: {str(e)}")
1123
1124@app.get("/case/{case_id}")
1125async def get_case_details(case_id: str):
1126 """Get detailed information about a specific case"""
1127 try:
1128 case_documents = []
1129
1130 for doc_id, doc_data in processor.local_documents.items():
1131 if doc_data.get('case_id', '').strip() == case_id.strip():
1132 case_documents.append({
1133 'doc_id': doc_id,
1134 'case_id': doc_data.get('case_id'),
1135 'device': doc_data.get('device'),
1136 'examiner': doc_data.get('examiner'),
1137 'date': doc_data.get('date'),
1138 'artifacts': doc_data.get('content_parts', [])
1139 })
1140
1141 if not case_documents:
1142 raise HTTPException(status_code=404, detail=f"Case {case_id} not found")
1143
1144 # Aggregate statistics
1145 total_artifacts = sum(len(doc['artifacts']) for doc in case_documents)
1146 artifact_types = {}
1147
1148 for doc in case_documents:
1149 for artifact in doc['artifacts']:
1150 art_type = artifact.get('type', 'unknown')
1151 artifact_types[art_type] = artifact_types.get(art_type, 0) + 1
1152
1153 return {
1154 "case_id": case_id,
1155 "documents": case_documents,
1156 "total_artifacts": total_artifacts,
1157 "artifact_breakdown": artifact_types
1158 }
1159
1160 except HTTPException:
1161 raise
1162 except Exception as e:
1163 raise HTTPException(status_code=500, detail=f"Error retrieving case: {str(e)}")
1164
1165@app.post("/reindex")
1166async def trigger_reindex():
1167 """Manually trigger a full reindex"""
1168 try:
1169 processor._sync_from_mongodb()
1170 processor._rebuild_index()
1171 return {
1172 "status": "success",
1173 "message": "Reindexing completed",
1174 "documents": len(processor.local_documents),
1175 "embeddings": sum(len(embs) for embs in processor.local_embeddings.values())
1176 }
1177 except Exception as e:
1178 raise HTTPException(status_code=500, detail=f"Reindexing failed: {str(e)}")
1179
1180@app.get("/anomalies")
1181async def check_all_anomalies():
1182 """Check for anomalies across all cases or specific case"""
1183 try:
1184 all_anomalies = {}
1185
1186 for doc_id, doc_data in processor.local_documents.items():
1187 case_id = doc_data.get('case_id', 'unknown')
1188 anomalies_in_case = []
1189
1190 # Check content parts for anomalies
1191 for part in doc_data.get('content_parts', []):
1192 if part.get('type') == 'anomaly':
1193 anomalies_in_case.append({
1194 'content': part['content'],
1195 'metadata': part.get('metadata', {})
1196 })
1197
1198 if anomalies_in_case:
1199 if case_id not in all_anomalies:
1200 all_anomalies[case_id] = {