backup_system.py

#!/usr/bin/env python3"""Automated Backup System for Ransomware PreventionThis script implements a robust backup system following the 3-2-1 backup strategyto protect against ransomware attacks. It includes versioning, integrity checking,and automated backup scheduling.""" import osimport sysimport shutilimport hashlibimport jsonimport timeimport threadingimport argparseimport scheduleimport loggingfrom datetime import datetime, timedeltafrom pathlib import Pathimport pickleimport zlibimport base64from colorama import init, Fore, Style init(autoreset=True) class BackupSystem:    """    Automated Backup System with 3-2-1 Strategy Implementation        3-2-1 Backup Strategy:    - 3 copies of data    - 2 different media types    - 1 copy offsite    """        def __init__(self, source_dirs=None, backup_root="/backup",                  retention_days=30, min_backup_interval=3600,                 integrity_check=True, compression=True):        """        Initialize Backup System                Args:            source_dirs: List of directories to backup            backup_root: Root directory for backups            retention_days: Number of days to retain backups            min_backup_interval: Minimum seconds between backups            integrity_check: Enable integrity checking            compression: Enable compression        """        self.source_dirs = source_dirs or [            str(Path.home() / "Documents"),            str(Path.home() / "Pictures"),            str(Path.home() / "Desktop"),            str(Path.home() / "Downloads")        ]        self.backup_root = backup_root        self.retention_days = retention_days        self.min_backup_interval = min_backup_interval        self.integrity_check = integrity_check        self.compression = compression                # Create backup directories        self.local_backup = os.path.join(backup_root, "local")        self.external_backup = os.path.join(backup_root, "external")        self.cloud_backup = os.path.join(backup_root, "cloud_staging")                for directory in [self.local_backup, self.external_backup, self.cloud_backup]:            os.makedirs(directory, exist_ok=True)                # Backup tracking        self.backup_manifest = os.path.join(backup_root, "backup_manifest.json")        self.last_backup_time = {}        self.backup_history = []                # Statistics        self.stats = {            'total_backups': 0,            'total_files': 0,            'total_size': 0,            'failed_backups': 0,            'integrity_failures': 0,            'start_time': time.time()        }                # Setup logging        self.setup_logging()                # Load existing manifest        self.load_manifest()        def setup_logging(self):        """Setup logging configuration"""        log_file = os.path.join(self.backup_root, 'backup.log')        logging.basicConfig(            level=logging.INFO,            format='%(asctime)s - %(levelname)s - %(message)s',            handlers=[                logging.FileHandler(log_file),                logging.StreamHandler()            ]        )        self.logger = logging.getLogger(__name__)        def load_manifest(self):        """Load backup manifest"""        if os.path.exists(self.backup_manifest):            try:                with open(self.backup_manifest, 'r') as f:                    data = json.load(f)                    self.last_backup_time = data.get('last_backup_time', {})                    self.backup_history = data.get('backup_history', [])                    self.stats.update(data.get('stats', {}))                self.logger.info(f"Loaded manifest with {len(self.backup_history)} backup records")            except Exception as e:                self.logger.error(f"Error loading manifest: {e}")        def save_manifest(self):        """Save backup manifest"""        try:            data = {                'last_backup_time': self.last_backup_time,                'backup_history': self.backup_history[-100:],  # Keep last 100                'stats': self.stats,                'updated': datetime.now().isoformat()            }            with open(self.backup_manifest, 'w') as f:                json.dump(data, f, indent=2)            self.logger.info("Manifest saved successfully")        except Exception as e:            self.logger.error(f"Error saving manifest: {e}")        def calculate_file_hash(self, filepath):        """Calculate SHA-256 hash of file"""        sha256 = hashlib.sha256()        try:            with open(filepath, 'rb') as f:                for chunk in iter(lambda: f.read(8192), b''):                    sha256.update(chunk)            return sha256.hexdigest()        except Exception as e:            self.logger.error(f"Error hashing {filepath}: {e}")            return None        def should_backup(self, source_dir):        """Check if directory should be backed up"""        # Check minimum interval        last_backup = self.last_backup_time.get(source_dir, 0)        if time.time() - last_backup < self.min_backup_interval:            return False                # Check if directory exists and has files        if not os.path.exists(source_dir):            return False                # Check for recent changes        recent_changes = self.check_recent_changes(source_dir)        return recent_changes        def check_recent_changes(self, directory):        """Check if files have been modified since last backup"""        last_backup = self.last_backup_time.get(directory, 0)                for root, dirs, files in os.walk(directory):            for file in files:                filepath = os.path.join(root, file)                try:                    mtime = os.path.getmtime(filepath)                    if mtime > last_backup:                        return True                except:                    continue                return False        def compress_file(self, filepath):        """Compress file using zlib"""        try:            with open(filepath, 'rb') as f:                data = f.read()                        compressed = zlib.compress(data, level=6)                        # Save compressed version            compressed_path = filepath + '.zlib'            with open(compressed_path, 'wb') as f:                f.write(compressed)                        return compressed_path        except Exception as e:            self.logger.error(f"Compression error for {filepath}: {e}")            return filepath        def backup_directory(self, source_dir, dest_dir):        """Backup a single directory"""        if not os.path.exists(source_dir):            self.logger.warning(f"Source directory does not exist: {source_dir}")            return 0, 0                # Create destination path        rel_path = os.path.basename(source_dir)        dest_path = os.path.join(dest_dir, rel_path)        os.makedirs(dest_path, exist_ok=True)                # Create backup info file        backup_info = {            'source': source_dir,            'destination': dest_path,            'timestamp': datetime.now().isoformat(),            'files': [],            'total_size': 0        }                files_backed_up = 0        total_size = 0                # Walk through source directory        for root, dirs, files in os.walk(source_dir):            for file in files:                try:                    source_file = os.path.join(root, file)                                        # Get relative path                    rel_file_path = os.path.relpath(source_file, source_dir)                    dest_file = os.path.join(dest_path, rel_file_path)                                        # Create destination directory                    os.makedirs(os.path.dirname(dest_file), exist_ok=True)                                        # Copy file with metadata                    if self.compression and file.endswith(('.txt', '.doc', '.docx', '.pdf', '.xls', '.xlsx')):                        # Compress text/document files                        temp_file = self.compress_file(source_file)                        shutil.copy2(temp_file, dest_file + '.compressed')                        os.remove(temp_file)                                                file_size = os.path.getsize(source_file)                        backup_info['files'].append({                            'name': file,                            'path': rel_file_path,                            'size': file_size,                            'compressed': True,                            'hash': self.calculate_file_hash(source_file)                        })                    else:                        # Normal copy for binary files                        shutil.copy2(source_file, dest_file)                                                file_size = os.path.getsize(source_file)                        backup_info['files'].append({                            'name': file,                            'path': rel_file_path,                            'size': file_size,                            'compressed': False,                            'hash': self.calculate_file_hash(source_file)                        })                                        files_backed_up += 1                    total_size += file_size                                        # Update stats                    self.stats['total_files'] += 1                    self.stats['total_size'] += file_size                                    except Exception as e:                    self.logger.error(f"Error backing up {file}: {e}")                    self.stats['failed_backups'] += 1                # Save backup info        info_file = os.path.join(dest_path, f"backup_info_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")        with open(info_file, 'w') as f:            json.dump(backup_info, f, indent=2)                return files_backed_up, total_size        def _backup_content_hash(self, path, compressed):        """SHA-256 of the original content of a backed-up file."""        if not compressed:            return self.calculate_file_hash(path)        try:            with open(path, 'rb') as handle:                return hashlib.sha256(zlib.decompress(handle.read())).hexdigest()        except (OSError, zlib.error) as exc:            self.logger.error(f"Cannot read compressed backup {path}: {exc}")            return None     def verify_backup(self, backup_dir):        """Verify backup integrity"""        self.logger.info(f"Verifying backup in {backup_dir}")                # Find backup info files        info_files = []        for root, dirs, files in os.walk(backup_dir):            for file in files:                if file.startswith('backup_info_') and file.endswith('.json'):                    info_files.append(os.path.join(root, file))                verified = 0        failed = 0                for info_file in info_files:            try:                with open(info_file, 'r') as f:                    backup_info = json.load(f)                                # Verify each file                for file_info in backup_info.get('files', []):                    # Files are stored next to their backup_info file, not directly under backup_dir.                    backup_path = os.path.join(os.path.dirname(info_file), file_info['path'])                    compressed = file_info.get('compressed')                    if compressed:                        backup_path += '.compressed'                                        if os.path.exists(backup_path):                        # The recorded hash is of the ORIGINAL content, so compressed copies                        # must be decompressed before hashing.                        if file_info.get('hash'):                            current_hash = self._backup_content_hash(backup_path, compressed)                            if current_hash != file_info['hash']:                                self.logger.error(f"Integrity check failed for {file_info['path']}")                                self.stats['integrity_failures'] += 1                                failed += 1                            else:                                verified += 1                    else:                        self.logger.error(f"Backup file missing: {file_info['path']}")                        failed += 1                            except Exception as e:                self.logger.error(f"Error verifying {info_file}: {e}")                failed += 1                self.logger.info(f"Verification complete: {verified} files OK, {failed} files failed")        return verified, failed        def rotate_backups(self, backup_dir):        """Remove old backups based on retention policy"""        self.logger.info(f"Rotating backups in {backup_dir}")                cutoff_time = time.time() - (self.retention_days * 24 * 3600)        removed = 0                for item in os.listdir(backup_dir):            item_path = os.path.join(backup_dir, item)                        if os.path.isdir(item_path):                # Check directory modification time                mtime = os.path.getmtime(item_path)                if mtime < cutoff_time:                    try:                        shutil.rmtree(item_path)                        self.logger.info(f"Removed old backup: {item}")                        removed += 1                    except Exception as e:                        self.logger.error(f"Error removing {item_path}: {e}")                return removed        def create_incremental_backup(self):        """Create incremental backup (only changed files)"""        self.logger.info("Starting incremental backup...")                timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')        backup_id = f"incremental_{timestamp}"                total_files = 0        total_size = 0                for source_dir in self.source_dirs:            if self.should_backup(source_dir):                # Determine backup destination                dest_dir = os.path.join(self.local_backup, backup_id)                                files, size = self.backup_directory(source_dir, dest_dir)                total_files += files                total_size += size                                # Update last backup time                self.last_backup_time[source_dir] = time.time()                # Record backup        backup_record = {            'id': backup_id,            'timestamp': timestamp,            'type': 'incremental',            'files': total_files,            'size': total_size,            'source_dirs': self.source_dirs        }                self.backup_history.append(backup_record)        self.stats['total_backups'] += 1                self.logger.info(f"Incremental backup complete: {total_files} files, {self.format_size(total_size)}")        self.save_manifest()                return backup_record        def create_full_backup(self):        """Create full backup of all directories"""        self.logger.info("Starting full backup...")                timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')        backup_id = f"full_{timestamp}"                total_files = 0        total_size = 0                # Backup to local storage        local_dest = os.path.join(self.local_backup, backup_id)        for source_dir in self.source_dirs:            files, size = self.backup_directory(source_dir, local_dest)            total_files += files            total_size += size                # Create copy for external backup (simulated)        if os.path.exists('/media/backup') or os.path.exists('/mnt/backup'):            external_dest = os.path.join(self.external_backup, backup_id)            self.sync_to_external(local_dest, external_dest)                # Stage for cloud backup        cloud_dest = os.path.join(self.cloud_backup, backup_id)        self.stage_for_cloud(local_dest, cloud_dest)                # Update last backup time        for source_dir in self.source_dirs:            self.last_backup_time[source_dir] = time.time()                # Record backup        backup_record = {            'id': backup_id,            'timestamp': timestamp,            'type': 'full',            'files': total_files,            'size': total_size,            'source_dirs': self.source_dirs        }                self.backup_history.append(backup_record)        self.stats['total_backups'] += 1                self.logger.info(f"Full backup complete: {total_files} files, {self.format_size(total_size)}")        self.save_manifest()                # Verify backup        if self.integrity_check:            verified, failed = self.verify_backup(local_dest)            self.logger.info(f"Backup verification: {verified} verified, {failed} failed")                return backup_record        def sync_to_external(self, source, destination):        """Sync backup to external drive"""        self.logger.info(f"Syncing to external storage: {destination}")                try:            os.makedirs(destination, exist_ok=True)                        # Use rsync if available, otherwise copy            if shutil.which('rsync'):                os.system(f"rsync -av --delete {source}/ {destination}/")            else:                # Fallback to shutil                for item in os.listdir(source):                    s = os.path.join(source, item)                    d = os.path.join(destination, item)                    if os.path.isdir(s):                        shutil.copytree(s, d, dirs_exist_ok=True)                    else:                        shutil.copy2(s, d)                        self.logger.info("External sync complete")        except Exception as e:            self.logger.error(f"External sync failed: {e}")        def stage_for_cloud(self, source, destination):        """Stage backup for cloud upload"""        self.logger.info(f"Staging for cloud backup: {destination}")                try:            os.makedirs(destination, exist_ok=True)                        # Create compressed archive for cloud upload            archive_name = f"cloud_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.tar.gz"            archive_path = os.path.join(destination, archive_name)                        import tarfile            with tarfile.open(archive_path, "w:gz") as tar:                tar.add(source, arcname=os.path.basename(source))                        self.logger.info(f"Cloud staging complete: {archive_name}")        except Exception as e:            self.logger.error(f"Cloud staging failed: {e}")        def restore_backup(self, backup_id, restore_path):        """Restore from backup"""        self.logger.info(f"Restoring backup {backup_id} to {restore_path}")                # Find backup        backup_found = None        for record in self.backup_history:            if record['id'] == backup_id:                backup_found = record                break                if not backup_found:            self.logger.error(f"Backup {backup_id} not found")            return False                # Locate backup files        backup_locations = [            os.path.join(self.local_backup, backup_id),            os.path.join(self.external_backup, backup_id),            os.path.join(self.cloud_backup, backup_id)        ]                restored = 0        for location in backup_locations:            if os.path.exists(location):                self.logger.info(f"Restoring from {location}")                                for root, dirs, files in os.walk(location):                    for file in files:                        if file.startswith('backup_info_'):                            continue                                                source_file = os.path.join(root, file)                        rel_path = os.path.relpath(source_file, location)                        dest_file = os.path.join(restore_path, rel_path)                                                os.makedirs(os.path.dirname(dest_file), exist_ok=True)                        shutil.copy2(source_file, dest_file)                        restored += 1                self.logger.info(f"Restore complete: {restored} files restored")        return True        def list_backups(self):        """List all available backups"""        print(f"\n{Fore.CYAN}📋 Available Backups:{Style.RESET_ALL}")        print(f"{Fore.CYAN}{'='*60}{Style.RESET_ALL}")                for record in sorted(self.backup_history, key=lambda x: x['timestamp'], reverse=True):            status = "✓" if record.get('verified') else " "            print(f"[{status}] {record['timestamp']} - {record['type'].upper()}")            print(f"    ID: {record['id']}")            print(f"    Files: {record['files']} | Size: {self.format_size(record['size'])}")            print()        def format_size(self, size):        """Format file size"""        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:            if size < 1024.0:                return f"{size:.1f} {unit}"            size /= 1024.0        return f"{size:.1f} PB"        def scheduled_backup(self):        """Run scheduled backup"""        self.logger.info("Running scheduled backup")                # Check if we need full backup (weekly)        last_full = None        for record in reversed(self.backup_history):            if record['type'] == 'full':                last_full = record                break                if not last_full or (time.time() - datetime.fromisoformat(last_full['timestamp']).timestamp()) > 7*24*3600:            self.create_full_backup()        else:            self.create_incremental_backup()                # Rotate old backups        removed = self.rotate_backups(self.local_backup)        if removed > 0:            self.logger.info(f"Removed {removed} old backups")        def monitor_ransomware_activity(self):        """Monitor for ransomware-like activity and trigger emergency backup"""        # This would integrate with detection scripts        # For now, just check for rapid file changes        try:            import psutil                        # Check for high disk I/O            disk_io = psutil.disk_io_counters()            if disk_io:                # Simple heuristic: if write bytes > threshold, trigger backup                if disk_io.write_bytes > 100 * 1024 * 1024:  # 100MB                    self.logger.warning("High disk write activity detected - triggering emergency backup")                    self.create_incremental_backup()                            except Exception as e:            self.logger.error(f"Error monitoring ransomware: {e}")        def start_automated_backup(self, interval_hours=6):        """Start automated backup scheduler"""        self.logger.info(f"Starting automated backup every {interval_hours} hours")                # Schedule regular backups        schedule.every(interval_hours).hours.do(self.scheduled_backup)                # Schedule daily verification        schedule.every().day.at("02:00").do(lambda: self.verify_backup(self.local_backup))                # Schedule weekly full backup        schedule.every().sunday.at("03:00").do(self.create_full_backup)                # Monitor for ransomware (every 5 minutes)        schedule.every(5).minutes.do(self.monitor_ransomware_activity)                print(f"{Fore.GREEN}[✓] Automated backup scheduler started{Style.RESET_ALL}")        print(f"{Fore.YELLOW}   Press Ctrl+C to stop{Style.RESET_ALL}")                try:            while True:                schedule.run_pending()                time.sleep(60)        except KeyboardInterrupt:            self.logger.info("Automated backup stopped")        def print_status(self):        """Print backup system status"""        uptime = time.time() - self.stats['start_time']                print(f"\n{Fore.CYAN}{'='*60}{Style.RESET_ALL}")        print(f"{Fore.CYAN}📊 BACKUP SYSTEM STATUS{Style.RESET_ALL}")        print(f"{Fore.CYAN}{'='*60}{Style.RESET_ALL}")        print(f"  Uptime: {int(uptime)} seconds")        print(f"  Total Backups: {self.stats['total_backups']}")        print(f"  Total Files: {self.stats['total_files']}")        print(f"  Total Size: {self.format_size(self.stats['total_size'])}")        print(f"  Failed Backups: {self.stats['failed_backups']}")        print(f"  Integrity Failures: {self.stats['integrity_failures']}")        print(f"{Fore.CYAN}{'='*60}{Style.RESET_ALL}")                print(f"\n{Fore.GREEN}Backup Locations:{Style.RESET_ALL}")        print(f"  Local: {self.local_backup}")        print(f"  External: {self.external_backup}")        print(f"  Cloud Staging: {self.cloud_backup}")                # Show last backup        if self.backup_history:            last = self.backup_history[-1]            print(f"\n{Fore.YELLOW}Last Backup:{Style.RESET_ALL}")            print(f"  Time: {last['timestamp']}")            print(f"  Type: {last['type'].upper()}")            print(f"  Files: {last['files']}")            print(f"  Size: {self.format_size(last['size'])}") def main():    parser = argparse.ArgumentParser(description='Automated Backup System for Ransomware Prevention')    parser.add_argument('--backup-root', default='/backup', help='Root directory for backups')    parser.add_argument('--source', help='Source directories to backup (comma-separated)')    parser.add_argument('--retention', type=int, default=30, help='Backup retention days')    parser.add_argument('--interval', type=int, default=6, help='Backup interval in hours')    parser.add_argument('--action', choices=['full', 'incremental', 'list', 'status', 'auto'],                       default='status', help='Action to perform')    parser.add_argument('--restore', help='Backup ID to restore')    parser.add_argument('--restore-path', help='Path to restore to')        args = parser.parse_args()        # Parse source directories    source_dirs = args.source.split(',') if args.source else None        # Initialize backup system    backup = BackupSystem(        source_dirs=source_dirs,        backup_root=args.backup_root,        retention_days=args.retention    )        if args.action == 'full':        backup.create_full_backup()    elif args.action == 'incremental':        backup.create_incremental_backup()    elif args.action == 'list':        backup.list_backups()    elif args.action == 'status':        backup.print_status()    elif args.action == 'auto':        backup.start_automated_backup(args.interval)        if args.restore and args.restore_path:        backup.restore_backup(args.restore, args.restore_path) if __name__ == "__main__":    main()