from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import HTMLResponse, JSONResponse import shutil import uvicorn import subprocess import os import json from typing import List, Dict import logging app = FastAPI(title="Disk Space Monitor", description="A simple API to monitor and manage disk space") # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def get_disk_space(): """Get disk space information""" total, used, free = shutil.disk_usage("/") return { "total_gb": round(total / (1024**3), 2), "used_gb": round(used / (1024**3), 2), "free_gb": round(free / (1024**3), 2), "used_percentage": round((used / total) * 100, 2) } def get_available_disks() -> List[Dict]: """Get list of available disks""" try: result = subprocess.run(['lsblk', '-J', '-o', 'NAME,SIZE,TYPE,MOUNTPOINT'], capture_output=True, text=True, check=True) disks_data = json.loads(result.stdout) disks = [] for device in disks_data['blockdevices']: if device['type'] == 'disk': disks.append({ 'name': device['name'], 'size': device['size'], 'type': device['type'], 'mountpoint': device.get('mountpoint', '') }) return disks except Exception as e: logger.error(f"Error getting disks: {e}") return [] def format_disk(disk_name: str, filesystem: str = "ext4") -> Dict: """ Format a disk with specified filesystem WARNING: This will erase all data on the disk! """ try: # Validate disk exists disks = get_available_disks() disk_exists = any(disk['name'] == disk_name for disk in disks) if not disk_exists: raise HTTPException(status_code=404, detail=f"Disk {disk_name} not found") # Check if disk is mounted for disk in disks: if disk['name'] == disk_name and disk.get('mountpoint'): raise HTTPException(status_code=400, detail=f"Disk {disk_name} is mounted. Unmount first.") # Format the disk (WARNING: DESTRUCTIVE OPERATION) device_path = f"/dev/{disk_name}" # Unmount if somehow still mounted subprocess.run(['umount', device_path], capture_output=True) # Create partition table (GPT) subprocess.run(['parted', '-s', device_path, 'mklabel', 'gpt'], check=True) # Create single partition subprocess.run(['parted', '-s', device_path, 'mkpart', 'primary', filesystem, '0%', '100%'], check=True) # Format the partition partition_path = f"{device_path}1" if filesystem == "ext4": subprocess.run(['mkfs.ext4', '-F', partition_path], check=True) elif filesystem == "xfs": subprocess.run(['mkfs.xfs', '-f', partition_path], check=True) elif filesystem == "ntfs": subprocess.run(['mkfs.ntfs', '-F', partition_path], check=True) else: raise HTTPException(status_code=400, detail=f"Unsupported filesystem: {filesystem}") return { "status": "success", "message": f"Disk {disk_name} formatted with {filesystem} filesystem", "disk": disk_name, "filesystem": filesystem } except subprocess.CalledProcessError as e: logger.error(f"Formatting error: {e}") raise HTTPException(status_code=500, detail=f"Formatting failed: {str(e)}") except Exception as e: logger.error(f"Unexpected error: {e}") raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}") @app.get("/") async def root(): """Root endpoint with HTML interface""" disk_info = get_disk_space() available_disks = get_available_disks() disks_html = "" for disk in available_disks: disks_html += f"""
{disk['name']} {disk['size']} {disk.get('mountpoint', 'Not mounted')}
""" html_content = f""" Disk Space Monitor

🖥️ Disk Space Monitor & Manager

📊 Current Disk Usage

Total Space: {disk_info['total_gb']} GB
Used Space: {disk_info['used_gb']} GB
Free Space: {disk_info['free_gb']} GB
Usage: {disk_info['used_percentage']}%

💾 Available Disks

{disks_html if disks_html else "

No disks found

"}

⚠️ Format Disk (DANGER ZONE)

⚠️ WARNING: Formatting will erase ALL data on the disk! This operation cannot be undone!

API Endpoints: /api/disk-space | /api/disks | /docs

""" return HTMLResponse(content=html_content) @app.get("/api/disk-space") async def get_disk_space_api(): """API endpoint to get disk space information in JSON format""" return get_disk_space() @app.get("/api/disks") async def get_disks_api(): """API endpoint to get list of available disks""" return {"disks": get_available_disks()} @app.post("/api/format-disk") async def format_disk_api(disk: str, filesystem: str = "ext4"): """ API endpoint to format a disk WARNING: This will erase all data on the disk! """ return format_disk(disk, filesystem) @app.get("/health") async def health_check(): """Health check endpoint""" return {"status": "healthy", "service": "disk-space-monitor"} if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)