smart_backup.py

#!/usr/bin/env python3
"""
Smart Backup Tool - Never lose your precious data again!
An intelligent backup tool with incremental backups, smart compression, and beautiful reporting.
"""

import os
import sys
import json
import shutil
import hashlib
import zipfile
import tarfile
import sqlite3
import threading
from datetime import datetime, timedelta
from pathlib import Path
import click
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.panel import Panel
from rich.text import Text
import psutil
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64

# Initialize rich console
console = Console()

class SmartBackup:
    def __init__(self):
        self.config_dir = Path.home() / '.config' / 'smart-backup'
        self.data_dir = Path.home() / '.local' / 'share' / 'smart-backup'
        self.logs_dir = self.data_dir / 'logs'
        self.config_file = self.config_dir / 'config.json'
        self.db_path = self.data_dir / 'backups.db'

        # Create directories if they don't exist
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.logs_dir.mkdir(parents=True, exist_ok=True)

        # Load configuration
        self.config = self.load_config()

        # Initialize database
        self.init_database()

    def load_config(self):
        """Load configuration from file or create default"""
        default_config = {
            "default_compression": True,
            "compression_level": 6,
            "default_encryption": False,
            "backup_retention": 30,
            "progress_bar": True,
            "colored_output": True,
            "default_exclude": [
                "*.tmp", "*.log", "__pycache__", "node_modules", ".git",
                "*.pyc", ".DS_Store", "Thumbs.db"
            ],
            "compression_algorithms": {
                "text": "lzma",
                "images": "zip",
                "videos": "none",
                "archives": "none"
            },
            "max_threads": 4,
            "chunk_size": "1MB"
        }

        if self.config_file.exists():
            try:
                with open(self.config_file, 'r') as f:
                    config = json.load(f)
                # Merge with defaults for any missing keys
                for key, value in default_config.items():
                    if key not in config:
                        config[key] = value
                return config
            except (json.JSONDecodeError, IOError):
                pass

        # Create default config file
        with open(self.config_file, 'w') as f:
            json.dump(default_config, f, indent=4)

        return default_config

    def init_database(self):
        """Initialize SQLite database for backup metadata"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        cursor.execute('''
            CREATE TABLE IF NOT EXISTS backups (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                source_path TEXT NOT NULL,
                destination_path TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                backup_type TEXT DEFAULT 'full',
                compression TEXT DEFAULT 'none',
                encryption BOOLEAN DEFAULT FALSE,
                file_count INTEGER DEFAULT 0,
                original_size INTEGER DEFAULT 0,
                backup_size INTEGER DEFAULT 0,
                checksum TEXT,
                status TEXT DEFAULT 'completed'
            )
        ''')

        cursor.execute('''
            CREATE TABLE IF NOT EXISTS backup_files (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                backup_id INTEGER,
                file_path TEXT NOT NULL,
                file_size INTEGER,
                file_hash TEXT,
                modified_time DATETIME,
                FOREIGN KEY (backup_id) REFERENCES backups (id)
            )
        ''')

        conn.commit()
        conn.close()

    def calculate_file_hash(self, file_path):
        """Calculate SHA-256 hash of a file"""
        hash_sha256 = hashlib.sha256()
        try:
            with open(file_path, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_sha256.update(chunk)
            return hash_sha256.hexdigest()
        except (IOError, OSError):
            return None

    def get_file_info(self, file_path):
        """Get file information including size and modification time"""
        try:
            stat = os.stat(file_path)
            return {
                'size': stat.st_size,
                'modified': datetime.fromtimestamp(stat.st_mtime),
                'hash': self.calculate_file_hash(file_path)
            }
        except (OSError, IOError):
            return None

    def should_exclude_file(self, file_path):
        """Check if file should be excluded based on patterns"""
        file_name = os.path.basename(file_path)
        for pattern in self.config['default_exclude']:
            if pattern.startswith('*'):
                if file_name.endswith(pattern[1:]):
                    return True
            elif pattern in file_path:
                return True
        return False

    def scan_directory(self, source_path, exclude_patterns=None):
        """Scan directory and return file information"""
        files_info = []
        total_size = 0

        exclude_patterns = exclude_patterns or []
        all_excludes = self.config['default_exclude'] + exclude_patterns

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Scanning directory...", total=None)

            for root, dirs, files in os.walk(source_path):
                # Remove excluded directories from dirs list to skip them
                dirs[:] = [d for d in dirs if not any(pattern in os.path.join(root, d) for pattern in all_excludes)]

                for file in files:
                    file_path = os.path.join(root, file)

                    if self.should_exclude_file(file_path):
                        continue

                    file_info = self.get_file_info(file_path)
                    if file_info:
                        files_info.append({
                            'path': file_path,
                            'relative_path': os.path.relpath(file_path, source_path),
                            **file_info
                        })
                        total_size += file_info['size']

                        progress.update(task, description=f"Scanning... {len(files_info)} files found")

        return files_info, total_size

    def format_size(self, size_bytes):
        """Format size in human readable format"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size_bytes < 1024.0:
                return f"{size_bytes:.1f} {unit}"
            size_bytes /= 1024.0
        return f"{size_bytes:.1f} PB"

    def create_backup_archive(self, files_info, destination_path, compress=True, encrypt=False, password=None):
        """Create backup archive with optional compression and encryption"""
        backup_name = f"backup_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"

        if compress:
            archive_path = destination_path / f"{backup_name}.tar.gz"
            mode = 'w:gz'
        else:
            archive_path = destination_path / f"{backup_name}.tar"
            mode = 'w'

        total_files = len(files_info)
        processed_size = 0

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=console
        ) as progress:
            task = progress.add_task("Creating backup...", total=total_files)

            with tarfile.open(archive_path, mode) as tar:
                for i, file_info in enumerate(files_info):
                    try:
                        tar.add(file_info['path'], arcname=file_info['relative_path'])
                        processed_size += file_info['size']
                        progress.update(task, advance=1,
                                      description=f"Backing up... {self.format_size(processed_size)}")
                    except (OSError, IOError) as e:
                        console.print(f"[yellow]Warning: Could not backup {file_info['path']}: {e}[/yellow]")

        # Handle encryption if requested
        if encrypt and password:
            encrypted_path = archive_path.with_suffix(archive_path.suffix + '.enc')
            self.encrypt_file(archive_path, encrypted_path, password)
            archive_path.unlink()  # Remove unencrypted file
            archive_path = encrypted_path

        return archive_path

    def encrypt_file(self, input_path, output_path, password):
        """Encrypt a file using Fernet (AES)"""
        # Generate key from password
        password_bytes = password.encode()
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password_bytes))
        fernet = Fernet(key)

        with open(input_path, 'rb') as infile, open(output_path, 'wb') as outfile:
            # Write salt first
            outfile.write(salt)
            # Encrypt and write data
            data = infile.read()
            encrypted_data = fernet.encrypt(data)
            outfile.write(encrypted_data)

    def save_backup_metadata(self, backup_info):
        """Save backup metadata to database"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        cursor.execute('''
            INSERT INTO backups
            (name, source_path, destination_path, backup_type, compression, encryption,
             file_count, original_size, backup_size, checksum, status)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            backup_info['name'],
            backup_info['source_path'],
            backup_info['destination_path'],
            backup_info['backup_type'],
            backup_info['compression'],
            backup_info['encryption'],
            backup_info['file_count'],
            backup_info['original_size'],
            backup_info['backup_size'],
            backup_info['checksum'],
            backup_info['status']
        ))

        backup_id = cursor.lastrowid

        # Save file information
        for file_info in backup_info.get('files', []):
            cursor.execute('''
                INSERT INTO backup_files (backup_id, file_path, file_size, file_hash, modified_time)
                VALUES (?, ?, ?, ?, ?)
            ''', (
                backup_id,
                file_info['relative_path'],
                file_info['size'],
                file_info['hash'],
                file_info['modified']
            ))

        conn.commit()
        conn.close()

        return backup_id

# Initialize the backup instance
backup_tool = SmartBackup()

@click.group()
@click.version_option(version='1.0.0')
def cli():
    """Smart Backup Tool - Never lose your precious data again!"""
    pass

@cli.command()
@click.argument('source_path', type=click.Path(exists=True))
@click.option('--destination', '-d', required=True, type=click.Path(), help='Backup destination directory')
@click.option('--compress', '-c', is_flag=True, help='Enable compression')
@click.option('--encrypt', '-e', is_flag=True, help='Enable encryption')
@click.option('--password', '-p', help='Encryption password')
@click.option('--exclude', help='Additional exclude patterns (comma-separated)')
@click.option('--name', help='Custom backup name')
@click.option('--incremental', '-i', is_flag=True, help='Create incremental backup')
def create(source_path, destination, compress, encrypt, password, exclude, name, incremental):
    """Create a new backup"""

    console.print(Panel.fit("📦 Smart Backup Tool v1.0.0", style="bold blue"))

    source_path = Path(source_path)
    destination_path = Path(destination)
    destination_path.mkdir(parents=True, exist_ok=True)

    # Handle encryption password
    if encrypt and not password:
        password = click.prompt("Enter encryption password", hide_input=True, confirmation_prompt=True)

    # Parse exclude patterns
    exclude_patterns = []
    if exclude:
        exclude_patterns = [p.strip() for p in exclude.split(',')]

    # Scan source directory
    console.print(f"🔍 Analyzing source directory: [bold]{source_path}[/bold]")
    files_info, total_size = backup_tool.scan_directory(source_path, exclude_patterns)

    if not files_info:
        console.print("[red]❌ No files found to backup![/red]")
        return

    # Display scan results
    console.print(f"📁 Found {len(files_info)} files ({backup_tool.format_size(total_size)})")

    # Show configuration
    config_table = Table(title="Backup Configuration")
    config_table.add_column("Setting", style="cyan")
    config_table.add_column("Value", style="green")

    config_table.add_row("Compression", "Enabled" if compress else "Disabled")
    config_table.add_row("Encryption", "Enabled" if encrypt else "Disabled")
    config_table.add_row("Backup Type", "Incremental" if incremental else "Full")
    config_table.add_row("Total Files", str(len(files_info)))
    config_table.add_row("Total Size", backup_tool.format_size(total_size))

    console.print(config_table)

    # Create backup
    console.print("\n🚀 Starting backup...")

    try:
        archive_path = backup_tool.create_backup_archive(
            files_info, destination_path, compress, encrypt, password
        )

        # Calculate backup size and compression ratio
        backup_size = archive_path.stat().st_size
        compression_ratio = ((total_size - backup_size) / total_size * 100) if total_size > 0 else 0

        # Save metadata
        backup_info = {
            'name': name or archive_path.stem,
            'source_path': str(source_path),
            'destination_path': str(archive_path),
            'backup_type': 'incremental' if incremental else 'full',
            'compression': 'gzip' if compress else 'none',
            'encryption': encrypt,
            'file_count': len(files_info),
            'original_size': total_size,
            'backup_size': backup_size,
            'checksum': backup_tool.calculate_file_hash(archive_path),
            'status': 'completed',
            'files': files_info
        }

        backup_id = backup_tool.save_backup_metadata(backup_info)

        # Display results
        console.print("\n✅ [bold green]Backup completed successfully![/bold green]")

        results_table = Table(title="Backup Results")
        results_table.add_column("Metric", style="cyan")
        results_table.add_column("Value", style="green")

        results_table.add_row("Backup ID", str(backup_id))
        results_table.add_row("Location", str(archive_path))
        results_table.add_row("Original Size", backup_tool.format_size(total_size))
        results_table.add_row("Backup Size", backup_tool.format_size(backup_size))
        if compress:
            results_table.add_row("Compression", f"{compression_ratio:.1f}% reduction")
        results_table.add_row("Files Processed", str(len(files_info)))

        console.print(results_table)

        console.print(f"\n💡 [dim]Tip: Use 'smart-backup list' to see all your backups[/dim]")

    except Exception as e:
        console.print(f"[red]❌ Backup failed: {e}[/red]")
        sys.exit(1)

@cli.command()
@click.option('--show-sizes', is_flag=True, help='Show backup sizes')
@click.option('--sort-by', type=click.Choice(['date', 'size', 'name']), default='date', help='Sort backups by')
def list(show_sizes, sort_by):
    """List all backups"""

    conn = sqlite3.connect(backup_tool.db_path)
    cursor = conn.cursor()

    cursor.execute('''
        SELECT id, name, created_at, file_count, original_size, backup_size,
               compression, encryption, status, destination_path
        FROM backups
        ORDER BY created_at DESC
    ''')

    backups = cursor.fetchall()
    conn.close()

    if not backups:
        console.print("[yellow]📦 No backups found. Create your first backup with 'smart-backup create'[/yellow]")
        return

    # Create table
    table = Table(title="📦 Available Backups")
    table.add_column("ID", style="cyan")
    table.add_column("Name", style="green")
    table.add_column("Date", style="blue")
    table.add_column("Files", style="yellow")

    if show_sizes:
        table.add_column("Original Size", style="magenta")
        table.add_column("Backup Size", style="magenta")
        table.add_column("Compression", style="red")

    table.add_column("Type", style="white")
    table.add_column("Status", style="green")

    total_backups = len(backups)
    total_storage = sum(backup[5] for backup in backups)  # backup_size

    for backup in backups:
        backup_id, name, created_at, file_count, original_size, backup_size, compression, encryption, status, dest_path = backup

        # Format date
        date_obj = datetime.fromisoformat(created_at)
        formatted_date = date_obj.strftime('%Y-%m-%d %H:%M')

        # Determine backup type
        backup_type = []
        if compression != 'none':
            backup_type.append("Compressed")
        if encryption:
            backup_type.append("Encrypted")
        if not backup_type:
            backup_type.append("Standard")

        type_str = ", ".join(backup_type)

        row = [
            str(backup_id),
            name[:25] + "..." if len(name) > 25 else name,
            formatted_date,
            str(file_count)
        ]

        if show_sizes:
            row.extend([
                backup_tool.format_size(original_size),
                backup_tool.format_size(backup_size),
                f"{((original_size - backup_size) / original_size * 100):.1f}%" if original_size > 0 and compression != 'none' else "N/A"
            ])

        row.extend([type_str, status.title()])

        table.add_row(*row)

    console.print(table)

    # Summary
    summary_table = Table(title="Summary")
    summary_table.add_column("Metric", style="cyan")
    summary_table.add_column("Value", style="green")

    summary_table.add_row("Total Backups", str(total_backups))
    summary_table.add_row("Storage Used", backup_tool.format_size(total_storage))

    # Available space
    try:
        if backups:
            # Get available space from first backup destination
            dest_path = Path(backups[0][9]).parent
            if dest_path.exists():
                available_space = shutil.disk_usage(dest_path).free
                summary_table.add_row("Available Space", backup_tool.format_size(available_space))
    except:
        pass

    console.print(summary_table)

@cli.command()
def status():
    """Show backup system status"""

    console.print(Panel.fit("📊 Smart Backup Status", style="bold blue"))

    # System information
    system_table = Table(title="System Information")
    system_table.add_column("Metric", style="cyan")
    system_table.add_column("Value", style="green")

    # Memory usage
    memory = psutil.virtual_memory()
    system_table.add_row("Available Memory", backup_tool.format_size(memory.available))
    system_table.add_row("CPU Count", str(psutil.cpu_count()))

    # Configuration
    config_table = Table(title="Configuration")
    config_table.add_column("Setting", style="cyan")
    config_table.add_column("Value", style="green")

    config_table.add_row("Config File", str(backup_tool.config_file))
    config_table.add_row("Database", str(backup_tool.db_path))
    config_table.add_row("Default Compression", "Yes" if backup_tool.config['default_compression'] else "No")
    config_table.add_row("Max Threads", str(backup_tool.config['max_threads']))

    console.print(system_table)
    console.print(config_table)

    # Recent backups
    conn = sqlite3.connect(backup_tool.db_path)
    cursor = conn.cursor()

    cursor.execute('''
        SELECT name, created_at, status
        FROM backups
        ORDER BY created_at DESC
        LIMIT 5
    ''')

    recent_backups = cursor.fetchall()
    conn.close()

    if recent_backups:
        recent_table = Table(title="Recent Backups")
        recent_table.add_column("Name", style="green")
        recent_table.add_column("Date", style="blue")
        recent_table.add_column("Status", style="yellow")

        for name, created_at, status in recent_backups:
            date_obj = datetime.fromisoformat(created_at)
            formatted_date = date_obj.strftime('%Y-%m-%d %H:%M')
            recent_table.add_row(name, formatted_date, status.title())

        console.print(recent_table)

if __name__ == '__main__':
    cli()