core/device_scanner.py

"""
Device Scanner Module
=====================

Scans Android device storage for recoverable data.
Identifies deleted files, databases, and media content.
"""

import os
import logging
import hashlib
from typing import List, Dict, Optional, Callable, Any
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed

from .adb_manager import ADBManager

logger = logging.getLogger(__name__)


class FileType(Enum):
    """Supported file types for recovery"""
    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"
    DOCUMENT = "document"
    DATABASE = "database"
    CONTACT = "contact"
    MESSAGE = "message"
    CALL_LOG = "call_log"
    APP_DATA = "app_data"
    WHATSAPP = "whatsapp"
    OTHER = "other"


class FileStatus(Enum):
    """File recovery status"""
    EXISTING = "existing"
    DELETED = "deleted"
    PARTIAL = "partial"
    CORRUPTED = "corrupted"
    UNKNOWN = "unknown"


@dataclass
class ScannedFile:
    """Represents a scanned file on the device"""
    path: str
    name: str
    size: int
    file_type: FileType
    status: FileStatus
    modified_time: Optional[datetime] = None
    created_time: Optional[datetime] = None
    mime_type: str = ""
    extension: str = ""
    signature: str = ""
    is_recoverable: bool = True
    recovery_chance: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def size_human(self) -> str:
        """Return human-readable file size"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if self.size < 1024:
                return f"{self.size:.1f} {unit}"
            self.size /= 1024
        return f"{self.size:.1f} TB"


@dataclass
class ScanResult:
    """Results from a device scan"""
    total_files: int = 0
    recoverable_files: int = 0
    total_size: int = 0
    recoverable_size: int = 0
    files: List[ScannedFile] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    scan_time: float = 0.0
    scan_depth: int = 1
    
    def get_files_by_type(self, file_type: FileType) -> List[ScannedFile]:
        """Get all files of a specific type"""
        return [f for f in self.files if f.file_type == file_type]
    
    def get_recoverable_files(self) -> List[ScannedFile]:
        """Get all recoverable files"""
        return [f for f in self.files if f.is_recoverable]


class DeviceScanner:
    """
    Scans Android device for recoverable data.
    
    Supports multiple scan modes:
    - Quick scan: Scans common locations only
    - Normal scan: Scans all accessible storage
    - Deep scan: Includes deleted file recovery (requires root)
    """
    
    # File signatures for identification
    FILE_SIGNATURES = {
        # Images
        b'\xff\xd8\xff': (FileType.IMAGE, 'jpg'),
        b'\x89PNG\r\n\x1a\n': (FileType.IMAGE, 'png'),
        b'GIF87a': (FileType.IMAGE, 'gif'),
        b'GIF89a': (FileType.IMAGE, 'gif'),
        b'BM': (FileType.IMAGE, 'bmp'),
        b'RIFF': (FileType.IMAGE, 'webp'),  # Also used for other formats
        
        # Videos
        b'\x00\x00\x00\x18ftypmp4': (FileType.VIDEO, 'mp4'),
        b'\x00\x00\x00\x1cftypmp4': (FileType.VIDEO, 'mp4'),
        b'\x00\x00\x00\x14ftyp3gp': (FileType.VIDEO, '3gp'),
        b'\x1aE\xdf\xa3': (FileType.VIDEO, 'mkv'),
        
        # Audio
        b'ID3': (FileType.AUDIO, 'mp3'),
        b'\xff\xfb': (FileType.AUDIO, 'mp3'),
        b'\xff\xf3': (FileType.AUDIO, 'mp3'),
        b'OggS': (FileType.AUDIO, 'ogg'),
        
        # Documents
        b'%PDF': (FileType.DOCUMENT, 'pdf'),
        b'PK\x03\x04': (FileType.DOCUMENT, 'zip'),  # Also docx, xlsx, etc.
        b'\xd0\xcf\x11\xe0': (FileType.DOCUMENT, 'doc'),
        
        # Databases
        b'SQLite format 3': (FileType.DATABASE, 'db'),
    }
    
    # Common storage paths on Android
    STORAGE_PATHS = {
        'internal': '/sdcard',
        'dcim': '/sdcard/DCIM',
        'pictures': '/sdcard/Pictures',
        'downloads': '/sdcard/Download',
        'documents': '/sdcard/Documents',
        'music': '/sdcard/Music',
        'movies': '/sdcard/Movies',
        'whatsapp': '/sdcard/WhatsApp',
        'telegram': '/sdcard/Telegram',
    }
    
    # Database paths (require root)
    DATABASE_PATHS = {
        'contacts': '/data/data/com.android.providers.contacts/databases/',
        'messages': '/data/data/com.android.providers.telephony/databases/',
        'whatsapp_db': '/data/data/com.whatsapp/databases/',
    }
    
    def __init__(self, adb_manager: ADBManager):
        """
        Initialize Device Scanner.
        
        Args:
            adb_manager: ADB Manager instance for device communication
        """
        self.adb = adb_manager
        self._scan_cancelled = False
        self._progress_callback: Optional[Callable[[int, int, str], None]] = None
        self._current_scan: Optional[ScanResult] = None
        self._lock = threading.Lock()
    
    def set_progress_callback(self, callback: Callable[[int, int, str], None]):
        """
        Set callback for progress updates.
        
        Args:
            callback: Function(current, total, message) for progress updates
        """
        self._progress_callback = callback
    
    def _update_progress(self, current: int, total: int, message: str):
        """Update scan progress"""
        if self._progress_callback:
            self._progress_callback(current, total, message)
    
    def cancel_scan(self):
        """Cancel the current scan operation"""
        self._scan_cancelled = True
    
    def quick_scan(self) -> ScanResult:
        """
        Perform a quick scan of common storage locations.
        
        Returns:
            ScanResult with found files
        """
        return self._scan(depth=1)
    
    def normal_scan(self) -> ScanResult:
        """
        Perform a normal scan of all accessible storage.
        
        Returns:
            ScanResult with found files
        """
        return self._scan(depth=2)
    
    def deep_scan(self) -> ScanResult:
        """
        Perform a deep scan including deleted file recovery.
        Requires root access for full functionality.
        
        Returns:
            ScanResult with found files
        """
        return self._scan(depth=3)
    
    def _scan(self, depth: int = 2) -> ScanResult:
        """
        Internal scan implementation.
        
        Args:
            depth: Scan depth (1=quick, 2=normal, 3=deep)
            
        Returns:
            ScanResult with found files
        """
        import time
        start_time = time.time()
        
        self._scan_cancelled = False
        result = ScanResult(scan_depth=depth)
        self._current_scan = result
        
        device_info = self.adb.get_device_info()
        if not device_info:
            result.errors.append("No device connected")
            return result
        
        logger.info(f"Starting scan (depth={depth}) on device: {device_info.serial}")
        
        # Determine paths to scan based on depth
        paths_to_scan = []
        
        if depth >= 1:
            # Quick scan - common locations
            paths_to_scan.extend([
                self.STORAGE_PATHS['dcim'],
                self.STORAGE_PATHS['pictures'],
                self.STORAGE_PATHS['downloads'],
            ])
        
        if depth >= 2:
            # Normal scan - all storage
            paths_to_scan.extend([
                self.STORAGE_PATHS['internal'],
                self.STORAGE_PATHS['documents'],
                self.STORAGE_PATHS['music'],
                self.STORAGE_PATHS['movies'],
                self.STORAGE_PATHS['whatsapp'],
            ])
        
        if depth >= 3 and device_info.is_rooted:
            # Deep scan - include databases (requires root)
            paths_to_scan.extend([
                self.DATABASE_PATHS['contacts'],
                self.DATABASE_PATHS['messages'],
                self.DATABASE_PATHS['whatsapp_db'],
            ])
        
        # Remove duplicates while preserving order
        seen = set()
        unique_paths = []
        for path in paths_to_scan:
            if path not in seen:
                seen.add(path)
                unique_paths.append(path)
        
        total_paths = len(unique_paths)
        
        # Scan each path
        for i, path in enumerate(unique_paths):
            if self._scan_cancelled:
                logger.info("Scan cancelled by user")
                break
            
            self._update_progress(i + 1, total_paths, f"Scanning: {path}")
            
            try:
                files = self._scan_directory(path, recursive=True, 
                                            use_root=depth >= 3 and device_info.is_rooted)
                
                with self._lock:
                    result.files.extend(files)
                    result.total_files += len(files)
                    
                    for f in files:
                        result.total_size += f.size
                        if f.is_recoverable:
                            result.recoverable_files += 1
                            result.recoverable_size += f.size
                            
            except Exception as e:
                error_msg = f"Error scanning {path}: {str(e)}"
                logger.error(error_msg)
                result.errors.append(error_msg)
        
        result.scan_time = time.time() - start_time
        logger.info(f"Scan completed: {result.total_files} files found in {result.scan_time:.2f}s")
        
        return result
    
    def _scan_directory(self, path: str, recursive: bool = True,
                       use_root: bool = False) -> List[ScannedFile]:
        """
        Scan a directory for files.
        
        Args:
            path: Directory path to scan
            recursive: Scan subdirectories
            use_root: Use root privileges
            
        Returns:
            List of ScannedFile objects
        """
        files = []
        
        # Check if path exists
        if not self.adb.file_exists(path, as_root=use_root):
            return files
        
        # List directory contents
        entries = self.adb.list_directory(path, as_root=use_root)
        
        for entry in entries:
            if self._scan_cancelled:
                break
            
            name = entry['name']
            
            # Skip special entries
            if name in ['.', '..']:
                continue
            
            full_path = f"{path}/{name}"
            
            if entry['is_directory']:
                if recursive:
                    # Recursively scan subdirectory
                    subfiles = self._scan_directory(full_path, recursive, use_root)
                    files.extend(subfiles)
            else:
                # Process file
                scanned_file = self._process_file(full_path, entry, use_root)
                if scanned_file:
                    files.append(scanned_file)
        
        return files
    
    def _process_file(self, path: str, entry: Dict[str, Any],
                     use_root: bool = False) -> Optional[ScannedFile]:
        """
        Process a file and create ScannedFile object.
        
        Args:
            path: File path
            entry: Directory entry information
            use_root: Use root privileges
            
        Returns:
            ScannedFile object or None if file should be skipped
        """
        name = entry['name']
        size = entry['size']
        
        # Determine file type from extension
        extension = os.path.splitext(name)[1].lower().lstrip('.')
        file_type = self._get_file_type_from_extension(extension)
        
        # Try to get file signature for better identification
        signature = ""
        if size > 0 and size < 100 * 1024 * 1024:  # Skip very large files
            signature = self._get_file_signature(path, use_root)
            if signature:
                detected_type, detected_ext = self._identify_by_signature(signature)
                if detected_type:
                    file_type = detected_type
                    if not extension:
                        extension = detected_ext
        
        # Parse modification time
        modified_time = None
        try:
            # Try to parse date from ls output
            date_str = entry.get('date', '')
            if date_str:
                modified_time = datetime.strptime(date_str, "%Y-%m-%d %H:%M")
        except ValueError:
            pass
        
        # Determine MIME type
        mime_type = self._get_mime_type(extension)
        
        # Create ScannedFile object
        scanned_file = ScannedFile(
            path=path,
            name=name,
            size=size,
            file_type=file_type,
            status=FileStatus.EXISTING,
            modified_time=modified_time,
            mime_type=mime_type,
            extension=extension,
            signature=signature,
            is_recoverable=True,
            recovery_chance=1.0
        )
        
        return scanned_file
    
    def _get_file_signature(self, path: str, use_root: bool = False) -> str:
        """Get the first bytes of a file for signature identification."""
        # Read first 32 bytes
        success, output = self.adb.shell(
            f"head -c 32 '{path}' | xxd -p",
            as_root=use_root
        )
        
        if success:
            return output.strip().replace('\n', '')
        return ""
    
    def _identify_by_signature(self, signature: str) -> tuple:
        """
        Identify file type by signature.
        
        Args:
            signature: Hex string of file header
            
        Returns:
            Tuple of (FileType, extension) or (None, None)
        """
        try:
            sig_bytes = bytes.fromhex(signature[:32])  # First 16 bytes
        except ValueError:
            return None, None
        
        for sig, (file_type, ext) in self.FILE_SIGNATURES.items():
            if sig_bytes.startswith(sig):
                return file_type, ext
        
        return None, None
    
    def _get_file_type_from_extension(self, extension: str) -> FileType:
        """Determine file type from extension."""
        extension = extension.lower()
        
        image_exts = {'jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp', 'heic', 'heif'}
        video_exts = {'mp4', 'avi', 'mkv', 'mov', '3gp', 'webm', 'flv'}
        audio_exts = {'mp3', 'wav', 'ogg', 'flac', 'm4a', 'aac', 'wma'}
        doc_exts = {'pdf', 'doc', 'docx', 'xls', 'xlsx', 'ppt', 'pptx', 'txt', 'rtf'}
        db_exts = {'db', 'sqlite', 'sqlite3'}
        
        if extension in image_exts:
            return FileType.IMAGE
        elif extension in video_exts:
            return FileType.VIDEO
        elif extension in audio_exts:
            return FileType.AUDIO
        elif extension in doc_exts:
            return FileType.DOCUMENT
        elif extension in db_exts:
            return FileType.DATABASE
        else:
            return FileType.OTHER
    
    def _get_mime_type(self, extension: str) -> str:
        """Get MIME type from extension."""
        mime_types = {
            'jpg': 'image/jpeg',
            'jpeg': 'image/jpeg',
            'png': 'image/png',
            'gif': 'image/gif',
            'bmp': 'image/bmp',
            'webp': 'image/webp',
            'mp4': 'video/mp4',
            'avi': 'video/x-msvideo',
            'mkv': 'video/x-matroska',
            'mov': 'video/quicktime',
            '3gp': 'video/3gpp',
            'mp3': 'audio/mpeg',
            'wav': 'audio/wav',
            'ogg': 'audio/ogg',
            'flac': 'audio/flac',
            'pdf': 'application/pdf',
            'doc': 'application/msword',
            'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
            'xls': 'application/vnd.ms-excel',
            'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            'txt': 'text/plain',
            'db': 'application/x-sqlite3',
        }
        return mime_types.get(extension.lower(), 'application/octet-stream')
    
    def scan_for_deleted_files(self, partition: str = "/data") -> List[ScannedFile]:
        """
        Scan for deleted files using raw partition analysis.
        Requires root access.
        
        Args:
            partition: Partition to scan
            
        Returns:
            List of potentially recoverable deleted files
        """
        device_info = self.adb.get_device_info()
        if not device_info or not device_info.is_rooted:
            logger.warning("Root access required for deleted file recovery")
            return []
        
        deleted_files = []
        
        # This is a simplified implementation
        # Real implementation would involve:
        # 1. Reading raw partition data
        # 2. Searching for file signatures
        # 3. Carving files from unallocated space
        
        logger.info(f"Scanning partition {partition} for deleted files...")
        
        # Use debugfs or similar tools if available
        success, output = self.adb.shell(
            f"find {partition} -name '*.deleted' 2>/dev/null",
            as_root=True
        )
        
        if success and output.strip():
            for line in output.strip().split('\n'):
                if line:
                    deleted_files.append(ScannedFile(
                        path=line,
                        name=os.path.basename(line),
                        size=0,
                        file_type=FileType.OTHER,
                        status=FileStatus.DELETED,
                        is_recoverable=True,
                        recovery_chance=0.5
                    ))
        
        return deleted_files
    
    def get_scan_statistics(self) -> Dict[str, Any]:
        """Get statistics from the current/last scan."""
        if not self._current_scan:
            return {}
        
        result = self._current_scan
        
        # Count files by type
        type_counts = {}
        for file_type in FileType:
            files = result.get_files_by_type(file_type)
            if files:
                type_counts[file_type.value] = {
                    'count': len(files),
                    'size': sum(f.size for f in files)
                }
        
        return {
            'total_files': result.total_files,
            'recoverable_files': result.recoverable_files,
            'total_size': result.total_size,
            'recoverable_size': result.recoverable_size,
            'scan_time': result.scan_time,
            'scan_depth': result.scan_depth,
            'errors': len(result.errors),
            'by_type': type_counts
        }