app.py

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"""
        <div class="disk-item">
            <span class="disk-name">{disk['name']}</span>
            <span class="disk-size">{disk['size']}</span>
            <span class="disk-mount">{disk.get('mountpoint', 'Not mounted')}</span>
        </div>
        """
    
    html_content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>Disk Space Monitor</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 40px; background-color: #f5f5f5; }}
            .container {{ max-width: 800px; margin: 0 auto; background: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }}
            h1 {{ color: #333; text-align: center; }}
            .section {{ margin: 30px 0; }}
            .disk-info, .disks-list, .format-section {{ 
                background: #f8f9fa; padding: 20px; border-radius: 8px; margin: 20px 0; 
            }}
            .metric {{ display: flex; justify-content: space-between; margin: 10px 0; padding: 10px; background: white; border-radius: 5px; }}
            .metric-label {{ font-weight: bold; color: #555; }}
            .metric-value {{ color: #007bff; font-weight: bold; }}
            .progress-bar {{ width: 100%; height: 20px; background: #e9ecef; border-radius: 10px; overflow: hidden; margin: 10px 0; }}
            .progress-fill {{ height: 100%; background: linear-gradient(90deg, #28a745, #ffc107, #dc3545); }}
            .refresh-btn {{ background: #007bff; color: white; border: none; padding: 10px 20px; border-radius: 5px; cursor: pointer; margin: 10px 0; }}
            .refresh-btn:hover {{ background: #0056b3; }}
            .disk-item {{ display: flex; justify-content: space-between; padding: 8px; background: white; margin: 5px 0; border-radius: 5px; }}
            .warning {{ color: #dc3545; font-weight: bold; background: #ffeaea; padding: 15px; border-radius: 5px; margin: 10px 0; }}
            .format-form {{ margin: 15px 0; }}
            .format-form select, .format-form input {{ padding: 8px; margin: 5px; border: 1px solid #ddd; border-radius: 4px; }}
            .format-btn {{ background: #dc3545; color: white; border: none; padding: 10px 20px; border-radius: 5px; cursor: pointer; }}
            .format-btn:hover {{ background: #c82333; }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🖥️ Disk Space Monitor & Manager</h1>
            
            <div class="section">
                <h2>📊 Current Disk Usage</h2>
                <div class="disk-info">
                    <div class="metric">
                        <span class="metric-label">Total Space:</span>
                        <span class="metric-value">{disk_info['total_gb']} GB</span>
                    </div>
                    <div class="metric">
                        <span class="metric-label">Used Space:</span>
                        <span class="metric-value">{disk_info['used_gb']} GB</span>
                    </div>
                    <div class="metric">
                        <span class="metric-label">Free Space:</span>
                        <span class="metric-value">{disk_info['free_gb']} GB</span>
                    </div>
                    <div class="metric">
                        <span class="metric-label">Usage:</span>
                        <span class="metric-value">{disk_info['used_percentage']}%</span>
                    </div>
                    <div class="progress-bar">
                        <div class="progress-fill" style="width: {disk_info['used_percentage']}%"></div>
                    </div>
                    <button class="refresh-btn" onclick="location.reload()">🔄 Refresh</button>
                </div>
            </div>

            <div class="section">
                <h2>💾 Available Disks</h2>
                <div class="disks-list">
                    {disks_html if disks_html else "<p>No disks found</p>"}
                </div>
            </div>

            <div class="section">
                <h2>⚠️ Format Disk (DANGER ZONE)</h2>
                <div class="warning">
                    ⚠️ WARNING: Formatting will erase ALL data on the disk! This operation cannot be undone!
                </div>
                <div class="format-section">
                    <form class="format-form" onsubmit="formatDisk(event)">
                        <select id="diskSelect" required>
                            <option value="">Select a disk</option>
                            {"".join([f'<option value="{disk["name"]}">{disk["name"]} ({disk["size"]})</option>' for disk in available_disks])}
                        </select>
                        <select id="fsSelect" required>
                            <option value="ext4">ext4</option>
                            <option value="xfs">XFS</option>
                            <option value="ntfs">NTFS</option>
                        </select>
                        <button type="submit" class="format-btn">🚫 Format Disk</button>
                    </form>
                </div>
            </div>

            <p style="text-align: center; color: #666; margin-top: 30px;">
                API Endpoints: <a href="/api/disk-space">/api/disk-space</a> | 
                <a href="/api/disks">/api/disks</a> | 
                <a href="/docs">/docs</a>
            </p>
        </div>

        <script>
            async function formatDisk(event) {{
                event.preventDefault();
                const disk = document.getElementById('diskSelect').value;
                const filesystem = document.getElementById('fsSelect').value;
                
                if (!disk) {{
                    alert('Please select a disk');
                    return;
                }}
                
                if (!confirm(`⚠️ WARNING: This will ERASE ALL DATA on disk ${{disk}}! Are you absolutely sure?`)) {{
                    return;
                }}
                
                try {{
                    const response = await fetch('/api/format-disk', {{
                        method: 'POST',
                        headers: {{ 'Content-Type': 'application/json' }},
                        body: JSON.stringify({{ disk: disk, filesystem: filesystem }})
                    }});
                    
                    const result = await response.json();
                    if (response.ok) {{
                        alert('✅ ' + result.message);
                        location.reload();
                    }} else {{
                        alert('❌ ' + result.detail);
                    }}
                }} catch (error) {{
                    alert('❌ Error: ' + error.message);
                }}
            }}
        </script>
    </body>
    </html>
    """
    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)