Upload core/data_recovery.py with huggingface_hub

core/data_recovery.py

@@ -0,0 +1,692 @@
+"""
+Data Recovery Module
+====================
+
+Handles the actual recovery of files from Android devices.
+Supports various recovery methods including direct copy, database extraction,
+and file carving for deleted files.
+"""
+
+import os
+import logging
+import shutil
+import tempfile
+import hashlib
+from typing import List, Dict, Optional, Callable, Any, Tuple
+from dataclasses import dataclass, field
+from enum import Enum
+from datetime import datetime
+from pathlib import Path
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from .adb_manager import ADBManager
+from .device_scanner import ScannedFile, FileType, FileStatus, ScanResult
+
+logger = logging.getLogger(__name__)
+
+
+class RecoveryStatus(Enum):
+    """Status of a recovery operation"""
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    PARTIAL = "partial"
+
+
+@dataclass
+class RecoveryTask:
+    """Represents a single file recovery task"""
+    source_file: ScannedFile
+    destination: str
+    status: RecoveryStatus = RecoveryStatus.PENDING
+    progress: float = 0.0
+    error_message: str = ""
+    recovered_size: int = 0
+    checksum: str = ""
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    
+    @property
+    def duration(self) -> float:
+        """Get recovery duration in seconds"""
+        if self.start_time and self.end_time:
+            return (self.end_time - self.start_time).total_seconds()
+        return 0.0
+
+
+@dataclass
+class RecoveryResult:
+    """Results from a recovery operation"""
+    total_files: int = 0
+    recovered_files: int = 0
+    failed_files: int = 0
+    total_size: int = 0
+    recovered_size: int = 0
+    tasks: List[RecoveryTask] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    output_directory: str = ""
+    
+    @property
+    def success_rate(self) -> float:
+        """Calculate recovery success rate"""
+        if self.total_files == 0:
+            return 0.0
+        return (self.recovered_files / self.total_files) * 100
+    
+    @property
+    def duration(self) -> float:
+        """Get total recovery duration in seconds"""
+        if self.start_time and self.end_time:
+            return (self.end_time - self.start_time).total_seconds()
+        return 0.0
+
+
+class DataRecovery:
+    """
+    Handles data recovery from Android devices.
+    
+    Features:
+    - Direct file copy for existing files
+    - Database extraction and parsing
+    - File carving for deleted files
+    - Parallel recovery for improved speed
+    - Progress tracking and cancellation support
+    """
+    
+    def __init__(self, adb_manager: ADBManager, output_dir: str = "./recovered_data"):
+        """
+        Initialize Data Recovery.
+        
+        Args:
+            adb_manager: ADB Manager instance
+            output_dir: Directory for recovered files
+        """
+        self.adb = adb_manager
+        self.output_dir = output_dir
+        self._cancelled = False
+        self._progress_callback: Optional[Callable[[int, int, str, float], None]] = None
+        self._current_result: Optional[RecoveryResult] = None
+        self._lock = threading.Lock()
+        self._max_workers = 4
+        
+        # Create output directory
+        os.makedirs(output_dir, exist_ok=True)
+    
+    def set_progress_callback(self, callback: Callable[[int, int, str, float], None]):
+        """
+        Set callback for progress updates.
+        
+        Args:
+            callback: Function(current, total, filename, progress) for updates
+        """
+        self._progress_callback = callback
+    
+    def _update_progress(self, current: int, total: int, filename: str, progress: float):
+        """Update recovery progress"""
+        if self._progress_callback:
+            self._progress_callback(current, total, filename, progress)
+    
+    def cancel_recovery(self):
+        """Cancel the current recovery operation"""
+        self._cancelled = True
+    
+    def recover_files(self, files: List[ScannedFile], 
+                     organize_by_type: bool = True) -> RecoveryResult:
+        """
+        Recover multiple files from the device.
+        
+        Args:
+            files: List of files to recover
+            organize_by_type: Organize recovered files by type
+            
+        Returns:
+            RecoveryResult with recovery details
+        """
+        self._cancelled = False
+        result = RecoveryResult(
+            total_files=len(files),
+            total_size=sum(f.size for f in files),
+            output_directory=self.output_dir,
+            start_time=datetime.now()
+        )
+        self._current_result = result
+        
+        logger.info(f"Starting recovery of {len(files)} files")
+        
+        # Create recovery tasks
+        tasks = []
+        for file in files:
+            # Determine destination path
+            if organize_by_type:
+                type_dir = self._get_type_directory(file.file_type)
+                dest_dir = os.path.join(self.output_dir, type_dir)
+            else:
+                dest_dir = self.output_dir
+            
+            os.makedirs(dest_dir, exist_ok=True)
+            destination = os.path.join(dest_dir, file.name)
+            
+            # Handle duplicate filenames
+            destination = self._get_unique_filename(destination)
+            
+            task = RecoveryTask(
+                source_file=file,
+                destination=destination
+            )
+            tasks.append(task)
+            result.tasks.append(task)
+        
+        # Execute recovery tasks
+        for i, task in enumerate(tasks):
+            if self._cancelled:
+                task.status = RecoveryStatus.CANCELLED
+                continue
+            
+            self._update_progress(i + 1, len(tasks), task.source_file.name, 0.0)
+            
+            try:
+                success = self._recover_single_file(task)
+                
+                with self._lock:
+                    if success:
+                        result.recovered_files += 1
+                        result.recovered_size += task.recovered_size
+                    else:
+                        result.failed_files += 1
+                        if task.error_message:
+                            result.errors.append(
+                                f"{task.source_file.name}: {task.error_message}"
+                            )
+                            
+            except Exception as e:
+                task.status = RecoveryStatus.FAILED
+                task.error_message = str(e)
+                result.failed_files += 1
+                result.errors.append(f"{task.source_file.name}: {str(e)}")
+                logger.error(f"Recovery failed for {task.source_file.name}: {e}")
+            
+            self._update_progress(i + 1, len(tasks), task.source_file.name, 100.0)
+        
+        result.end_time = datetime.now()
+        
+        logger.info(
+            f"Recovery completed: {result.recovered_files}/{result.total_files} files "
+            f"({result.success_rate:.1f}% success rate)"
+        )
+        
+        return result
+    
+    def recover_files_parallel(self, files: List[ScannedFile],
+                              organize_by_type: bool = True,
+                              max_workers: int = 4) -> RecoveryResult:
+        """
+        Recover multiple files in parallel.
+        
+        Args:
+            files: List of files to recover
+            organize_by_type: Organize recovered files by type
+            max_workers: Maximum parallel workers
+            
+        Returns:
+            RecoveryResult with recovery details
+        """
+        self._cancelled = False
+        self._max_workers = max_workers
+        
+        result = RecoveryResult(
+            total_files=len(files),
+            total_size=sum(f.size for f in files),
+            output_directory=self.output_dir,
+            start_time=datetime.now()
+        )
+        self._current_result = result
+        
+        logger.info(f"Starting parallel recovery of {len(files)} files with {max_workers} workers")
+        
+        # Create recovery tasks
+        tasks = []
+        for file in files:
+            if organize_by_type:
+                type_dir = self._get_type_directory(file.file_type)
+                dest_dir = os.path.join(self.output_dir, type_dir)
+            else:
+                dest_dir = self.output_dir
+            
+            os.makedirs(dest_dir, exist_ok=True)
+            destination = os.path.join(dest_dir, file.name)
+            destination = self._get_unique_filename(destination)
+            
+            task = RecoveryTask(
+                source_file=file,
+                destination=destination
+            )
+            tasks.append(task)
+            result.tasks.append(task)
+        
+        # Execute in parallel
+        completed = 0
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            future_to_task = {
+                executor.submit(self._recover_single_file, task): task
+                for task in tasks
+            }
+            
+            for future in as_completed(future_to_task):
+                if self._cancelled:
+                    executor.shutdown(wait=False)
+                    break
+                
+                task = future_to_task[future]
+                completed += 1
+                
+                try:
+                    success = future.result()
+                    
+                    with self._lock:
+                        if success:
+                            result.recovered_files += 1
+                            result.recovered_size += task.recovered_size
+                        else:
+                            result.failed_files += 1
+                            if task.error_message:
+                                result.errors.append(
+                                    f"{task.source_file.name}: {task.error_message}"
+                                )
+                                
+                except Exception as e:
+                    task.status = RecoveryStatus.FAILED
+                    task.error_message = str(e)
+                    result.failed_files += 1
+                    result.errors.append(f"{task.source_file.name}: {str(e)}")
+                
+                self._update_progress(
+                    completed, len(tasks), task.source_file.name, 100.0
+                )
+        
+        result.end_time = datetime.now()
+        
+        logger.info(
+            f"Parallel recovery completed: {result.recovered_files}/{result.total_files} files "
+            f"in {result.duration:.2f}s"
+        )
+        
+        return result
+    
+    def _recover_single_file(self, task: RecoveryTask) -> bool:
+        """
+        Recover a single file.
+        
+        Args:
+            task: Recovery task
+            
+        Returns:
+            True if recovery successful
+        """
+        task.status = RecoveryStatus.IN_PROGRESS
+        task.start_time = datetime.now()
+        
+        source = task.source_file
+        
+        try:
+            if source.status == FileStatus.EXISTING:
+                # Direct copy for existing files
+                success = self._copy_file(source.path, task.destination)
+            elif source.status == FileStatus.DELETED:
+                # File carving for deleted files
+                success = self._carve_file(source, task.destination)
+            else:
+                # Try direct copy as fallback
+                success = self._copy_file(source.path, task.destination)
+            
+            if success:
+                # Verify recovered file
+                if os.path.exists(task.destination):
+                    task.recovered_size = os.path.getsize(task.destination)
+                    task.checksum = self._calculate_checksum(task.destination)
+                    task.status = RecoveryStatus.COMPLETED
+                    task.progress = 100.0
+                else:
+                    task.status = RecoveryStatus.FAILED
+                    task.error_message = "File not found after recovery"
+                    success = False
+            else:
+                task.status = RecoveryStatus.FAILED
+                if not task.error_message:
+                    task.error_message = "Recovery failed"
+                    
+        except Exception as e:
+            task.status = RecoveryStatus.FAILED
+            task.error_message = str(e)
+            success = False
+        
+        task.end_time = datetime.now()
+        return success
+    
+    def _copy_file(self, source_path: str, destination: str) -> bool:
+        """
+        Copy a file from device to local storage.
+        
+        Args:
+            source_path: Path on device
+            destination: Local destination path
+            
+        Returns:
+            True if successful
+        """
+        return self.adb.pull_file(source_path, destination)
+    
+    def _carve_file(self, source: ScannedFile, destination: str) -> bool:
+        """
+        Attempt to carve a deleted file from device storage.
+        
+        Args:
+            source: Source file information
+            destination: Local destination path
+            
+        Returns:
+            True if successful
+        """
+        # This is a simplified implementation
+        # Real file carving would involve:
+        # 1. Reading raw partition data
+        # 2. Searching for file signatures
+        # 3. Extracting file data based on structure
+        
+        device_info = self.adb.get_device_info()
+        if not device_info or not device_info.is_rooted:
+            logger.warning("Root access required for file carving")
+            return False
+        
+        # Try to recover using file path if still accessible
+        if self.adb.file_exists(source.path, as_root=True):
+            return self._copy_file_as_root(source.path, destination)
+        
+        return False
+    
+    def _copy_file_as_root(self, source_path: str, destination: str) -> bool:
+        """
+        Copy a file using root privileges.
+        
+        Args:
+            source_path: Path on device
+            destination: Local destination path
+            
+        Returns:
+            True if successful
+        """
+        # Create temporary accessible copy
+        temp_path = f"/sdcard/.recovery_temp_{os.path.basename(source_path)}"
+        
+        # Copy file to accessible location
+        success, _ = self.adb.shell(f"cp '{source_path}' '{temp_path}'", as_root=True)
+        
+        if not success:
+            return False
+        
+        # Pull the file
+        result = self.adb.pull_file(temp_path, destination)
+        
+        # Clean up temp file
+        self.adb.shell(f"rm '{temp_path}'", as_root=True)
+        
+        return result
+    
+    def _get_type_directory(self, file_type: FileType) -> str:
+        """Get directory name for file type."""
+        type_dirs = {
+            FileType.IMAGE: "Images",
+            FileType.VIDEO: "Videos",
+            FileType.AUDIO: "Audio",
+            FileType.DOCUMENT: "Documents",
+            FileType.DATABASE: "Databases",
+            FileType.CONTACT: "Contacts",
+            FileType.MESSAGE: "Messages",
+            FileType.CALL_LOG: "CallLogs",
+            FileType.APP_DATA: "AppData",
+            FileType.WHATSAPP: "WhatsApp",
+            FileType.OTHER: "Other",
+        }
+        return type_dirs.get(file_type, "Other")
+    
+    def _get_unique_filename(self, path: str) -> str:
+        """
+        Get a unique filename if file already exists.
+        
+        Args:
+            path: Original file path
+            
+        Returns:
+            Unique file path
+        """
+        if not os.path.exists(path):
+            return path
+        
+        directory = os.path.dirname(path)
+        filename = os.path.basename(path)
+        name, ext = os.path.splitext(filename)
+        
+        counter = 1
+        while True:
+            new_path = os.path.join(directory, f"{name}_{counter}{ext}")
+            if not os.path.exists(new_path):
+                return new_path
+            counter += 1
+    
+    def _calculate_checksum(self, filepath: str, algorithm: str = "md5") -> str:
+        """
+        Calculate file checksum.
+        
+        Args:
+            filepath: Path to file
+            algorithm: Hash algorithm (md5, sha1, sha256)
+            
+        Returns:
+            Hex digest of checksum
+        """
+        if algorithm == "md5":
+            hasher = hashlib.md5()
+        elif algorithm == "sha1":
+            hasher = hashlib.sha1()
+        elif algorithm == "sha256":
+            hasher = hashlib.sha256()
+        else:
+            hasher = hashlib.md5()
+        
+        try:
+            with open(filepath, 'rb') as f:
+                for chunk in iter(lambda: f.read(8192), b''):
+                    hasher.update(chunk)
+            return hasher.hexdigest()
+        except Exception:
+            return ""
+    
+    def recover_contacts(self, output_file: str) -> bool:
+        """
+        Recover contacts from device database.
+        
+        Args:
+            output_file: Output file path (CSV or VCF)
+            
+        Returns:
+            True if successful
+        """
+        from .database_recovery import DatabaseRecovery
+        
+        db_recovery = DatabaseRecovery(self.adb)
+        contacts = db_recovery.extract_contacts()
+        
+        if not contacts:
+            return False
+        
+        # Export based on file extension
+        ext = os.path.splitext(output_file)[1].lower()
+        
+        if ext == '.vcf':
+            return self._export_contacts_vcf(contacts, output_file)
+        else:
+            return self._export_contacts_csv(contacts, output_file)
+    
+    def _export_contacts_csv(self, contacts: List[Dict], output_file: str) -> bool:
+        """Export contacts to CSV file."""
+        import csv
+        
+        try:
+            with open(output_file, 'w', newline='', encoding='utf-8') as f:
+                if contacts:
+                    writer = csv.DictWriter(f, fieldnames=contacts[0].keys())
+                    writer.writeheader()
+                    writer.writerows(contacts)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to export contacts to CSV: {e}")
+            return False
+    
+    def _export_contacts_vcf(self, contacts: List[Dict], output_file: str) -> bool:
+        """Export contacts to VCF (vCard) file."""
+        try:
+            with open(output_file, 'w', encoding='utf-8') as f:
+                for contact in contacts:
+                    f.write("BEGIN:VCARD\n")
+                    f.write("VERSION:3.0\n")
+                    
+                    name = contact.get('name', '')
+                    if name:
+                        f.write(f"FN:{name}\n")
+                        f.write(f"N:{name};;;;\n")
+                    
+                    phone = contact.get('phone', '')
+                    if phone:
+                        f.write(f"TEL;TYPE=CELL:{phone}\n")
+                    
+                    email = contact.get('email', '')
+                    if email:
+                        f.write(f"EMAIL:{email}\n")
+                    
+                    f.write("END:VCARD\n\n")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to export contacts to VCF: {e}")
+            return False
+    
+    def recover_messages(self, output_file: str) -> bool:
+        """
+        Recover SMS/MMS messages from device database.
+        
+        Args:
+            output_file: Output file path (CSV or XML)
+            
+        Returns:
+            True if successful
+        """
+        from .database_recovery import DatabaseRecovery
+        
+        db_recovery = DatabaseRecovery(self.adb)
+        messages = db_recovery.extract_messages()
+        
+        if not messages:
+            return False
+        
+        ext = os.path.splitext(output_file)[1].lower()
+        
+        if ext == '.xml':
+            return self._export_messages_xml(messages, output_file)
+        else:
+            return self._export_messages_csv(messages, output_file)
+    
+    def _export_messages_csv(self, messages: List[Dict], output_file: str) -> bool:
+        """Export messages to CSV file."""
+        import csv
+        
+        try:
+            with open(output_file, 'w', newline='', encoding='utf-8') as f:
+                if messages:
+                    writer = csv.DictWriter(f, fieldnames=messages[0].keys())
+                    writer.writeheader()
+                    writer.writerows(messages)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to export messages to CSV: {e}")
+            return False
+    
+    def _export_messages_xml(self, messages: List[Dict], output_file: str) -> bool:
+        """Export messages to XML file."""
+        try:
+            with open(output_file, 'w', encoding='utf-8') as f:
+                f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
+                f.write('<messages>\n')
+                
+                for msg in messages:
+                    f.write('  <message>\n')
+                    for key, value in msg.items():
+                        safe_value = str(value).replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
+                        f.write(f'    <{key}>{safe_value}</{key}>\n')
+                    f.write('  </message>\n')
+                
+                f.write('</messages>\n')
+            return True
+        except Exception as e:
+            logger.error(f"Failed to export messages to XML: {e}")
+            return False
+    
+    def recover_call_logs(self, output_file: str) -> bool:
+        """
+        Recover call logs from device database.
+        
+        Args:
+            output_file: Output file path
+            
+        Returns:
+            True if successful
+        """
+        from .database_recovery import DatabaseRecovery
+        
+        db_recovery = DatabaseRecovery(self.adb)
+        call_logs = db_recovery.extract_call_logs()
+        
+        if not call_logs:
+            return False
+        
+        return self._export_messages_csv(call_logs, output_file)
+    
+    def get_recovery_statistics(self) -> Dict[str, Any]:
+        """Get statistics from the current/last recovery."""
+        if not self._current_result:
+            return {}
+        
+        result = self._current_result
+        
+        # Count by status
+        status_counts = {}
+        for status in RecoveryStatus:
+            count = sum(1 for t in result.tasks if t.status == status)
+            if count > 0:
+                status_counts[status.value] = count
+        
+        # Count by file type
+        type_counts = {}
+        for task in result.tasks:
+            file_type = task.source_file.file_type.value
+            if file_type not in type_counts:
+                type_counts[file_type] = {'total': 0, 'recovered': 0}
+            type_counts[file_type]['total'] += 1
+            if task.status == RecoveryStatus.COMPLETED:
+                type_counts[file_type]['recovered'] += 1
+        
+        return {
+            'total_files': result.total_files,
+            'recovered_files': result.recovered_files,
+            'failed_files': result.failed_files,
+            'success_rate': result.success_rate,
+            'total_size': result.total_size,
+            'recovered_size': result.recovered_size,
+            'duration': result.duration,
+            'output_directory': result.output_directory,
+            'by_status': status_counts,
+            'by_type': type_counts,
+            'errors': result.errors[:10]  # First 10 errors
+        }