Upload core/database_recovery.py with huggingface_hub

core/database_recovery.py

@@ -0,0 +1,574 @@
+"""
+Database Recovery Module
+========================
+
+Handles extraction and parsing of Android SQLite databases.
+Supports contacts, messages, call logs, and app-specific databases.
+"""
+
+import os
+import logging
+import sqlite3
+import tempfile
+from typing import List, Dict, Optional, Any
+from datetime import datetime
+from pathlib import Path
+
+from .adb_manager import ADBManager
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseRecovery:
+    """
+    Extracts and parses data from Android SQLite databases.
+    
+    Supported databases:
+    - Contacts (contacts2.db)
+    - Messages (mmssms.db)
+    - Call Logs (calllog.db)
+    - WhatsApp (msgstore.db, wa.db)
+    """
+    
+    # Database paths on Android
+    DATABASE_PATHS = {
+        'contacts': '/data/data/com.android.providers.contacts/databases/contacts2.db',
+        'messages': '/data/data/com.android.providers.telephony/databases/mmssms.db',
+        'call_logs': '/data/data/com.android.providers.contacts/databases/calllog.db',
+        'whatsapp_messages': '/data/data/com.whatsapp/databases/msgstore.db',
+        'whatsapp_contacts': '/data/data/com.whatsapp/databases/wa.db',
+    }
+    
+    def __init__(self, adb_manager: ADBManager):
+        """
+        Initialize Database Recovery.
+        
+        Args:
+            adb_manager: ADB Manager instance
+        """
+        self.adb = adb_manager
+        self._temp_dir = tempfile.mkdtemp(prefix="android_recovery_")
+    
+    def __del__(self):
+        """Cleanup temporary files."""
+        import shutil
+        try:
+            shutil.rmtree(self._temp_dir, ignore_errors=True)
+        except Exception:
+            pass
+    
+    def _pull_database(self, db_path: str) -> Optional[str]:
+        """
+        Pull a database file from the device.
+        
+        Args:
+            db_path: Path to database on device
+            
+        Returns:
+            Local path to pulled database or None
+        """
+        device_info = self.adb.get_device_info()
+        if not device_info:
+            logger.error("No device connected")
+            return None
+        
+        # Check if database exists
+        if not self.adb.file_exists(db_path, as_root=True):
+            logger.warning(f"Database not found: {db_path}")
+            return None
+        
+        # Create local path
+        db_name = os.path.basename(db_path)
+        local_path = os.path.join(self._temp_dir, db_name)
+        
+        # Need root to access /data/data
+        if db_path.startswith('/data/data'):
+            if not device_info.is_rooted:
+                logger.warning("Root access required to access app databases")
+                return None
+            
+            # Copy to accessible location first
+            temp_device_path = f"/sdcard/.temp_{db_name}"
+            success, _ = self.adb.shell(f"cp '{db_path}' '{temp_device_path}'", as_root=True)
+            
+            if not success:
+                logger.error(f"Failed to copy database: {db_path}")
+                return None
+            
+            # Also copy journal and wal files if they exist
+            for suffix in ['-journal', '-wal', '-shm']:
+                journal_path = db_path + suffix
+                if self.adb.file_exists(journal_path, as_root=True):
+                    self.adb.shell(
+                        f"cp '{journal_path}' '{temp_device_path}{suffix}'",
+                        as_root=True
+                    )
+            
+            # Pull the database
+            if not self.adb.pull_file(temp_device_path, local_path):
+                logger.error(f"Failed to pull database: {temp_device_path}")
+                self.adb.shell(f"rm '{temp_device_path}'*", as_root=True)
+                return None
+            
+            # Pull journal files
+            for suffix in ['-journal', '-wal', '-shm']:
+                temp_journal = temp_device_path + suffix
+                if self.adb.file_exists(temp_journal):
+                    self.adb.pull_file(temp_journal, local_path + suffix)
+            
+            # Cleanup
+            self.adb.shell(f"rm '{temp_device_path}'*", as_root=True)
+        else:
+            # Direct pull for accessible paths
+            if not self.adb.pull_file(db_path, local_path):
+                logger.error(f"Failed to pull database: {db_path}")
+                return None
+        
+        return local_path
+    
+    def _query_database(self, db_path: str, query: str) -> List[Dict[str, Any]]:
+        """
+        Execute a query on a SQLite database.
+        
+        Args:
+            db_path: Path to local database file
+            query: SQL query to execute
+            
+        Returns:
+            List of result dictionaries
+        """
+        results = []
+        
+        try:
+            conn = sqlite3.connect(db_path)
+            conn.row_factory = sqlite3.Row
+            cursor = conn.cursor()
+            
+            cursor.execute(query)
+            rows = cursor.fetchall()
+            
+            for row in rows:
+                results.append(dict(row))
+            
+            conn.close()
+        except sqlite3.Error as e:
+            logger.error(f"Database query error: {e}")
+        except Exception as e:
+            logger.error(f"Error querying database: {e}")
+        
+        return results
+    
+    def extract_contacts(self) -> List[Dict[str, Any]]:
+        """
+        Extract contacts from the device.
+        
+        Returns:
+            List of contact dictionaries
+        """
+        db_path = self._pull_database(self.DATABASE_PATHS['contacts'])
+        if not db_path:
+            return []
+        
+        contacts = []
+        
+        try:
+            # Query for contacts with phone numbers
+            query = """
+                SELECT 
+                    c._id as id,
+                    c.display_name as name,
+                    p.data1 as phone,
+                    p.data2 as phone_type,
+                    e.data1 as email
+                FROM contacts c
+                LEFT JOIN data p ON c._id = p.contact_id AND p.mimetype_id = (
+                    SELECT _id FROM mimetypes WHERE mimetype = 'vnd.android.cursor.item/phone_v2'
+                )
+                LEFT JOIN data e ON c._id = e.contact_id AND e.mimetype_id = (
+                    SELECT _id FROM mimetypes WHERE mimetype = 'vnd.android.cursor.item/email_v2'
+                )
+                WHERE c.display_name IS NOT NULL
+                ORDER BY c.display_name
+            """
+            
+            results = self._query_database(db_path, query)
+            
+            # Process and deduplicate
+            seen_ids = set()
+            for row in results:
+                contact_id = row.get('id')
+                if contact_id in seen_ids:
+                    # Update existing contact with additional info
+                    for contact in contacts:
+                        if contact['id'] == contact_id:
+                            if row.get('phone') and row['phone'] not in contact.get('phones', []):
+                                contact.setdefault('phones', []).append(row['phone'])
+                            if row.get('email') and row['email'] not in contact.get('emails', []):
+                                contact.setdefault('emails', []).append(row['email'])
+                            break
+                else:
+                    seen_ids.add(contact_id)
+                    contact = {
+                        'id': contact_id,
+                        'name': row.get('name', ''),
+                        'phone': row.get('phone', ''),
+                        'email': row.get('email', ''),
+                        'phones': [row['phone']] if row.get('phone') else [],
+                        'emails': [row['email']] if row.get('email') else [],
+                    }
+                    contacts.append(contact)
+            
+            logger.info(f"Extracted {len(contacts)} contacts")
+            
+        except Exception as e:
+            logger.error(f"Error extracting contacts: {e}")
+        
+        return contacts
+    
+    def extract_messages(self) -> List[Dict[str, Any]]:
+        """
+        Extract SMS/MMS messages from the device.
+        
+        Returns:
+            List of message dictionaries
+        """
+        db_path = self._pull_database(self.DATABASE_PATHS['messages'])
+        if not db_path:
+            return []
+        
+        messages = []
+        
+        try:
+            # Query for SMS messages
+            query = """
+                SELECT 
+                    _id as id,
+                    address as phone_number,
+                    body as message,
+                    date as timestamp,
+                    type as message_type,
+                    read as is_read,
+                    seen as is_seen
+                FROM sms
+                ORDER BY date DESC
+            """
+            
+            results = self._query_database(db_path, query)
+            
+            for row in results:
+                # Convert timestamp
+                timestamp = row.get('timestamp', 0)
+                if timestamp:
+                    try:
+                        date_str = datetime.fromtimestamp(timestamp / 1000).strftime('%Y-%m-%d %H:%M:%S')
+                    except Exception:
+                        date_str = str(timestamp)
+                else:
+                    date_str = ''
+                
+                # Determine message direction
+                msg_type = row.get('message_type', 0)
+                direction = 'received' if msg_type == 1 else 'sent' if msg_type == 2 else 'unknown'
+                
+                message = {
+                    'id': row.get('id'),
+                    'phone_number': row.get('phone_number', ''),
+                    'message': row.get('message', ''),
+                    'date': date_str,
+                    'timestamp': timestamp,
+                    'direction': direction,
+                    'is_read': bool(row.get('is_read', 0)),
+                }
+                messages.append(message)
+            
+            logger.info(f"Extracted {len(messages)} messages")
+            
+        except Exception as e:
+            logger.error(f"Error extracting messages: {e}")
+        
+        return messages
+    
+    def extract_call_logs(self) -> List[Dict[str, Any]]:
+        """
+        Extract call logs from the device.
+        
+        Returns:
+            List of call log dictionaries
+        """
+        db_path = self._pull_database(self.DATABASE_PATHS['call_logs'])
+        if not db_path:
+            return []
+        
+        call_logs = []
+        
+        try:
+            query = """
+                SELECT 
+                    _id as id,
+                    number as phone_number,
+                    name as contact_name,
+                    date as timestamp,
+                    duration as duration_seconds,
+                    type as call_type
+                FROM calls
+                ORDER BY date DESC
+            """
+            
+            results = self._query_database(db_path, query)
+            
+            for row in results:
+                # Convert timestamp
+                timestamp = row.get('timestamp', 0)
+                if timestamp:
+                    try:
+                        date_str = datetime.fromtimestamp(timestamp / 1000).strftime('%Y-%m-%d %H:%M:%S')
+                    except Exception:
+                        date_str = str(timestamp)
+                else:
+                    date_str = ''
+                
+                # Determine call type
+                call_type_id = row.get('call_type', 0)
+                call_types = {
+                    1: 'incoming',
+                    2: 'outgoing',
+                    3: 'missed',
+                    4: 'voicemail',
+                    5: 'rejected',
+                    6: 'blocked',
+                }
+                call_type = call_types.get(call_type_id, 'unknown')
+                
+                # Format duration
+                duration = row.get('duration_seconds', 0)
+                if duration:
+                    minutes, seconds = divmod(duration, 60)
+                    hours, minutes = divmod(minutes, 60)
+                    if hours:
+                        duration_str = f"{hours}:{minutes:02d}:{seconds:02d}"
+                    else:
+                        duration_str = f"{minutes}:{seconds:02d}"
+                else:
+                    duration_str = "0:00"
+                
+                call_log = {
+                    'id': row.get('id'),
+                    'phone_number': row.get('phone_number', ''),
+                    'contact_name': row.get('contact_name', ''),
+                    'date': date_str,
+                    'timestamp': timestamp,
+                    'duration': duration_str,
+                    'duration_seconds': duration,
+                    'call_type': call_type,
+                }
+                call_logs.append(call_log)
+            
+            logger.info(f"Extracted {len(call_logs)} call logs")
+            
+        except Exception as e:
+            logger.error(f"Error extracting call logs: {e}")
+        
+        return call_logs
+    
+    def extract_whatsapp_messages(self) -> List[Dict[str, Any]]:
+        """
+        Extract WhatsApp messages from the device.
+        
+        Returns:
+            List of WhatsApp message dictionaries
+        """
+        db_path = self._pull_database(self.DATABASE_PATHS['whatsapp_messages'])
+        if not db_path:
+            return []
+        
+        messages = []
+        
+        try:
+            # WhatsApp database structure varies by version
+            # This is a common query structure
+            query = """
+                SELECT 
+                    m._id as id,
+                    m.key_remote_jid as chat_id,
+                    m.data as message,
+                    m.timestamp as timestamp,
+                    m.key_from_me as is_from_me,
+                    m.media_wa_type as media_type,
+                    m.media_mime_type as mime_type,
+                    m.media_name as media_name
+                FROM messages m
+                WHERE m.data IS NOT NULL OR m.media_wa_type > 0
+                ORDER BY m.timestamp DESC
+                LIMIT 10000
+            """
+            
+            results = self._query_database(db_path, query)
+            
+            for row in results:
+                # Convert timestamp
+                timestamp = row.get('timestamp', 0)
+                if timestamp:
+                    try:
+                        date_str = datetime.fromtimestamp(timestamp / 1000).strftime('%Y-%m-%d %H:%M:%S')
+                    except Exception:
+                        date_str = str(timestamp)
+                else:
+                    date_str = ''
+                
+                # Parse chat ID to get phone number
+                chat_id = row.get('chat_id', '')
+                phone_number = chat_id.split('@')[0] if '@' in chat_id else chat_id
+                
+                # Determine media type
+                media_types = {
+                    0: 'text',
+                    1: 'image',
+                    2: 'audio',
+                    3: 'video',
+                    4: 'contact',
+                    5: 'location',
+                    8: 'document',
+                    9: 'gif',
+                    13: 'sticker',
+                }
+                media_type = media_types.get(row.get('media_type', 0), 'unknown')
+                
+                message = {
+                    'id': row.get('id'),
+                    'phone_number': phone_number,
+                    'message': row.get('message', ''),
+                    'date': date_str,
+                    'timestamp': timestamp,
+                    'direction': 'sent' if row.get('is_from_me') else 'received',
+                    'media_type': media_type,
+                    'media_name': row.get('media_name', ''),
+                }
+                messages.append(message)
+            
+            logger.info(f"Extracted {len(messages)} WhatsApp messages")
+            
+        except Exception as e:
+            logger.error(f"Error extracting WhatsApp messages: {e}")
+        
+        return messages
+    
+    def extract_whatsapp_contacts(self) -> List[Dict[str, Any]]:
+        """
+        Extract WhatsApp contacts from the device.
+        
+        Returns:
+            List of WhatsApp contact dictionaries
+        """
+        db_path = self._pull_database(self.DATABASE_PATHS['whatsapp_contacts'])
+        if not db_path:
+            return []
+        
+        contacts = []
+        
+        try:
+            query = """
+                SELECT 
+                    _id as id,
+                    jid as whatsapp_id,
+                    display_name as name,
+                    number as phone_number,
+                    status as status_message
+                FROM wa_contacts
+                WHERE is_whatsapp_user = 1
+                ORDER BY display_name
+            """
+            
+            results = self._query_database(db_path, query)
+            
+            for row in results:
+                # Parse WhatsApp ID to get phone number if not available
+                wa_id = row.get('whatsapp_id', '')
+                phone = row.get('phone_number', '')
+                if not phone and wa_id:
+                    phone = wa_id.split('@')[0]
+                
+                contact = {
+                    'id': row.get('id'),
+                    'name': row.get('name', ''),
+                    'phone_number': phone,
+                    'whatsapp_id': wa_id,
+                    'status': row.get('status_message', ''),
+                }
+                contacts.append(contact)
+            
+            logger.info(f"Extracted {len(contacts)} WhatsApp contacts")
+            
+        except Exception as e:
+            logger.error(f"Error extracting WhatsApp contacts: {e}")
+        
+        return contacts
+    
+    def extract_all_data(self) -> Dict[str, List[Dict[str, Any]]]:
+        """
+        Extract all available data from the device.
+        
+        Returns:
+            Dictionary with all extracted data
+        """
+        return {
+            'contacts': self.extract_contacts(),
+            'messages': self.extract_messages(),
+            'call_logs': self.extract_call_logs(),
+            'whatsapp_messages': self.extract_whatsapp_messages(),
+            'whatsapp_contacts': self.extract_whatsapp_contacts(),
+        }
+    
+    def get_database_info(self, db_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Get information about a database.
+        
+        Args:
+            db_name: Database name (contacts, messages, etc.)
+            
+        Returns:
+            Dictionary with database information
+        """
+        if db_name not in self.DATABASE_PATHS:
+            return None
+        
+        db_path = self._pull_database(self.DATABASE_PATHS[db_name])
+        if not db_path:
+            return None
+        
+        info = {
+            'name': db_name,
+            'path': self.DATABASE_PATHS[db_name],
+            'local_path': db_path,
+            'size': os.path.getsize(db_path),
+            'tables': [],
+        }
+        
+        try:
+            conn = sqlite3.connect(db_path)
+            cursor = conn.cursor()
+            
+            # Get table list
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
+            tables = cursor.fetchall()
+            
+            for (table_name,) in tables:
+                # Get row count
+                cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
+                count = cursor.fetchone()[0]
+                
+                # Get column info
+                cursor.execute(f"PRAGMA table_info({table_name})")
+                columns = [col[1] for col in cursor.fetchall()]
+                
+                info['tables'].append({
+                    'name': table_name,
+                    'row_count': count,
+                    'columns': columns,
+                })
+            
+            conn.close()
+            
+        except Exception as e:
+            logger.error(f"Error getting database info: {e}")
+        
+        return info