default/telegram-agy-bot/backup_manager.py

import os
import sys
import time
import zipfile
import shutil
import logging
from pathlib import Path
from typing import List, Tuple, Optional, Any
from datetime import datetime

logger = logging.getLogger("AGYBackup")

# Telegram Bot API Document limit is 50MB (52,428,800 bytes).
# We use 48MB (50,331,648 bytes) per chunk for a safe margin.
CHUNK_SIZE_BYTES = 48 * 1024 * 1024

EXCLUDED_DIRS = {
    "__pycache__",
    "venv",
    ".venv",
    "node_modules",
    ".next",
    ".git",
    ".pytest_cache",
    ".cache",
    ".tmp",
    "dist",
    "build",
    "tmp",
    "uploads",
}

EXCLUDED_EXTENSIONS = {
    ".pyc",
    ".pyo",
    ".pyd",
    ".DS_Store",
}

def create_project_zip(workspace_path: Path, output_zip_path: Path) -> Tuple[int, int]:
    """
    Creates a zip archive of workspace_path at output_zip_path.
    Returns (total_files_count, total_uncompressed_bytes).
    """
    output_zip_path.parent.mkdir(parents=True, exist_ok=True)
    file_count = 0
    total_bytes = 0

    with zipfile.ZipFile(output_zip_path, "w", zipfile.ZIP_DEFLATED, compresslevel=6) as zip_file:
        for root, dirs, files in os.walk(workspace_path):
            # Prune excluded directories in-place
            dirs[:] = [d for d in dirs if d not in EXCLUDED_DIRS and not d.startswith(".")]

            for file_name in files:
                if any(file_name.endswith(ext) for ext in EXCLUDED_EXTENSIONS):
                    continue

                abs_file_path = Path(root) / file_name
                if not abs_file_path.is_file() or abs_file_path.is_symlink():
                    continue

                # Relative path inside the zip archive
                rel_path = abs_file_path.relative_to(workspace_path)
                
                try:
                    file_size = abs_file_path.stat().st_size
                    # Avoid archiving giant single files (>200MB) directly inside project backup if unwanted
                    if file_size > 300 * 1024 * 1024:
                        continue
                    zip_file.write(abs_file_path, arcname=str(rel_path))
                    file_count += 1
                    total_bytes += file_size
                except Exception as file_err:
                    logger.warning(f"Skipping file {abs_file_path} during backup: {file_err}")

    return file_count, total_bytes

def split_file_into_chunks(source_file: Path, chunk_size: int = CHUNK_SIZE_BYTES) -> List[Path]:
    """
    Splits source_file into chunks of chunk_size if larger than chunk_size.
    Returns list of Path objects for the resulting file(s).
    """
    total_size = source_file.stat().st_size
    if total_size <= chunk_size:
        return [source_file]

    chunks: List[Path] = []
    chunk_index = 1
    
    with open(source_file, "rb") as src_f:
        while True:
            chunk_data = src_f.read(chunk_size)
            if not chunk_data:
                break
            chunk_name = f"{source_file.name}.{chunk_index:03d}"
            chunk_path = source_file.parent / chunk_name
            with open(chunk_path, "wb") as chunk_f:
                chunk_f.write(chunk_data)
            chunks.append(chunk_path)
            chunk_index += 1

    return chunks

class BackupManager:
    def __init__(self, temp_dir: Path = Path("/tmp/agy_backups")):
        self.temp_dir = temp_dir
        self.temp_dir.mkdir(parents=True, exist_ok=True)

    def generate_backup(self, project_name: str, workspace_path: str) -> Tuple[List[Path], int, float, str]:
        """
        Creates zip archive of project workspace, splits it if necessary into <=48MB chunks.
        Returns: (chunk_paths, file_count, total_zip_size_mb, backup_session_id)
        """
        ws = Path(workspace_path).expanduser().resolve()
        if not ws.exists():
            raise FileNotFoundError(f"مسیر کاری پروژه یافت نشد: {workspace_path}")

        session_id = f"bk_{int(time.time())}_{os.getpid()}"
        session_dir = self.temp_dir / session_id
        session_dir.mkdir(parents=True, exist_ok=True)

        clean_proj_name = "".join(c for c in project_name if c.isalnum() or c in ("-", "_")).strip() or "project"
        timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
        zip_filename = f"{clean_proj_name}_backup_{timestamp_str}.zip"
        zip_path = session_dir / zip_filename

        logger.info(f"Creating project backup zip for '{project_name}' at {zip_path}...")
        file_count, uncompressed_bytes = create_project_zip(ws, zip_path)

        zip_size_bytes = zip_path.stat().st_size
        total_zip_size_mb = zip_size_bytes / (1024 * 1024)

        logger.info(f"Backup created: {file_count} files, zip size: {total_zip_size_mb:.2f} MB. Splitting if needed...")
        chunks = split_file_into_chunks(zip_path, CHUNK_SIZE_BYTES)

        return chunks, file_count, total_zip_size_mb, session_id

    def cleanup(self, session_id: str):
        """Cleans up temporary backup files after sending to save disk space."""
        session_dir = self.temp_dir / session_id
        if session_dir.exists():
            try:
                shutil.rmtree(session_dir, ignore_errors=True)
                logger.info(f"Cleaned up backup session directory: {session_dir}")
            except Exception as e:
                logger.error(f"Error cleaning up backup session {session_id}: {e}")

backup_manager = BackupManager()