"""Cleanup utilities for memory dump extractions and temporary files"""
import shutil
import time
import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
import zipfile
logger = logging.getLogger(__name__)
class ManagedExtraction:
"""Context manager for memory dump extraction with guaranteed cleanup"""
def __init__(self, dump_path: Path, extraction_dir: Path, dump_id: str,
auto_cleanup: bool = True):
"""
Initialize managed extraction
Args:
dump_path: Path to the dump file (possibly .zip)
extraction_dir: Base directory for extractions
dump_id: Identifier for this dump
auto_cleanup: Whether to cleanup on exit
"""
self.dump_path = dump_path
self.extraction_dir = extraction_dir
self.dump_id = dump_id
self.auto_cleanup = auto_cleanup
self.extracted_path = None
self.temp_dir = None
def __enter__(self) -> Path:
"""Extract the dump and return path to extracted file"""
timestamp = int(time.time())
self.temp_dir = self.extraction_dir / f"memdump_{self.dump_id}_{timestamp}"
self.temp_dir.mkdir(parents=True, exist_ok=True)
if self.dump_path.suffix.lower() == '.zip':
logger.info(f"Extracting {self.dump_path.name} to {self.temp_dir}")
with zipfile.ZipFile(self.dump_path, 'r') as zip_ref:
zip_ref.extractall(self.temp_dir)
# Find the actual memory dump file
for ext in ['.raw', '.mem', '.dmp', '.vmem', '.bin']:
extracted_files = list(self.temp_dir.rglob(f'*{ext}'))
if extracted_files:
self.extracted_path = extracted_files[0]
return self.extracted_path
# If no specific extension, return first file
files = list(self.temp_dir.iterdir())
if files:
self.extracted_path = files[0]
return self.extracted_path
raise ValueError(f"No memory dump found in {self.dump_path}")
else:
# Not a zip, return original path
self.extracted_path = self.dump_path
return self.extracted_path
def __exit__(self, exc_type, exc_val, exc_tb):
"""Cleanup extracted files"""
if self.auto_cleanup and self.temp_dir and self.temp_dir.exists():
try:
logger.info(f"Cleaning up extraction: {self.temp_dir}")
shutil.rmtree(self.temp_dir, ignore_errors=True)
except Exception as e:
logger.warning(f"Failed to cleanup {self.temp_dir}: {e}")
def cleanup_old_extractions(extraction_dir: Path,
retention_hours: int = 24,
dry_run: bool = False) -> Dict[str, Any]:
"""
Remove extraction directories older than retention period
Args:
extraction_dir: Directory containing extractions
retention_hours: Keep extractions newer than this many hours
dry_run: If True, don't actually delete, just report
Returns:
Dict with cleanup statistics
"""
if not extraction_dir.exists():
return {
'removed_count': 0,
'freed_bytes': 0,
'errors': []
}
cutoff_time = time.time() - (retention_hours * 3600)
removed_count = 0
freed_bytes = 0
errors = []
for extraction in extraction_dir.glob("memdump_*"):
if not extraction.is_dir():
continue
try:
mtime = extraction.stat().st_mtime
if mtime < cutoff_time:
# Calculate size before deletion
size = get_directory_size(extraction)
if dry_run:
logger.info(f"Would remove: {extraction} ({size / (1024**3):.2f} GB)")
else:
logger.info(f"Removing old extraction: {extraction} ({size / (1024**3):.2f} GB)")
shutil.rmtree(extraction)
removed_count += 1
freed_bytes += size
except Exception as e:
error_msg = f"Failed to remove {extraction}: {e}"
logger.error(error_msg)
errors.append(error_msg)
return {
'removed_count': removed_count,
'freed_bytes': freed_bytes,
'freed_gb': freed_bytes / (1024**3),
'errors': errors,
'dry_run': dry_run
}
def cleanup_all_extractions(extraction_dir: Path) -> Dict[str, Any]:
"""
Remove all extraction directories
Args:
extraction_dir: Directory containing extractions
Returns:
Dict with cleanup statistics
"""
if not extraction_dir.exists():
return {
'removed_count': 0,
'freed_bytes': 0,
'errors': []
}
removed_count = 0
freed_bytes = 0
errors = []
for extraction in extraction_dir.glob("memdump_*"):
if not extraction.is_dir():
continue
try:
size = get_directory_size(extraction)
logger.info(f"Removing extraction: {extraction} ({size / (1024**3):.2f} GB)")
shutil.rmtree(extraction)
removed_count += 1
freed_bytes += size
except Exception as e:
error_msg = f"Failed to remove {extraction}: {e}"
logger.error(error_msg)
errors.append(error_msg)
return {
'removed_count': removed_count,
'freed_bytes': freed_bytes,
'freed_gb': freed_bytes / (1024**3),
'errors': errors
}
def get_directory_size(path: Path) -> int:
"""
Calculate total size of a directory
Args:
path: Directory path
Returns:
Size in bytes
"""
total = 0
try:
for item in path.rglob('*'):
if item.is_file():
total += item.stat().st_size
except Exception as e:
logger.warning(f"Error calculating size for {path}: {e}")
return total
def get_disk_usage(data_dir: Path, extraction_dir: Path) -> Dict[str, Any]:
"""
Get disk usage statistics for MCP server
Args:
data_dir: Main data directory
extraction_dir: Extractions directory
Returns:
Dict with usage statistics
"""
stats = {
'data_dir': str(data_dir),
'extraction_dir': str(extraction_dir),
'database_size_bytes': 0,
'database_size_mb': 0,
'exports_size_bytes': 0,
'exports_size_mb': 0,
'extractions_size_bytes': 0,
'extractions_size_gb': 0,
'extractions_count': 0,
'total_size_bytes': 0,
'total_size_gb': 0
}
# Database size
db_file = data_dir / "artifacts.db"
if db_file.exists():
stats['database_size_bytes'] = db_file.stat().st_size
stats['database_size_mb'] = stats['database_size_bytes'] / (1024**2)
# Exports size
exports_dir = data_dir / "exports"
if exports_dir.exists():
stats['exports_size_bytes'] = get_directory_size(exports_dir)
stats['exports_size_mb'] = stats['exports_size_bytes'] / (1024**2)
# Extractions size
if extraction_dir.exists():
stats['extractions_size_bytes'] = get_directory_size(extraction_dir)
stats['extractions_size_gb'] = stats['extractions_size_bytes'] / (1024**3)
stats['extractions_count'] = len(list(extraction_dir.glob("memdump_*")))
# Total
stats['total_size_bytes'] = (
stats['database_size_bytes'] +
stats['exports_size_bytes'] +
stats['extractions_size_bytes']
)
stats['total_size_gb'] = stats['total_size_bytes'] / (1024**3)
return stats
def list_extractions(extraction_dir: Path) -> List[Dict[str, Any]]:
"""
List all current extractions with metadata
Args:
extraction_dir: Directory containing extractions
Returns:
List of extraction info dicts
"""
if not extraction_dir.exists():
return []
extractions = []
for extraction in extraction_dir.glob("memdump_*"):
if not extraction.is_dir():
continue
try:
stat = extraction.stat()
size = get_directory_size(extraction)
age_hours = (time.time() - stat.st_mtime) / 3600
extractions.append({
'path': str(extraction),
'name': extraction.name,
'size_bytes': size,
'size_gb': size / (1024**3),
'created': datetime.fromtimestamp(stat.st_mtime).isoformat(),
'age_hours': age_hours
})
except Exception as e:
logger.warning(f"Error reading {extraction}: {e}")
# Sort by age (newest first)
extractions.sort(key=lambda x: x['age_hours'])
return extractions
