ArdupilotCustomFirmwareBuilder/overlay_manager/main.py

import os
import sqlite3
import shutil
import tempfile
import logging
from urllib.parse import quote
from fastapi import FastAPI, Request, File, UploadFile, Form, HTTPException
from fastapi.responses import HTMLResponse, RedirectResponse, FileResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("overlay-manager")

app = FastAPI(root_path="/patch-manager")

os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

OVERLAY_DIR = os.path.abspath("/srv")

# Use the new mounted directory for the database
DB_DIR = "/app/db"
os.makedirs(DB_DIR, exist_ok=True)
DB_PATH = os.path.join(DB_DIR, "overlays.db")

def init_db():
    """Initializes the SQLite database and schema."""
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS profiles (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                is_active BOOLEAN NOT NULL DEFAULT 0
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS files (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                profile_id INTEGER NOT NULL,
                filepath TEXT NOT NULL,
                content BLOB NOT NULL,
                FOREIGN KEY(profile_id) REFERENCES profiles(id) ON DELETE CASCADE,
                UNIQUE(profile_id, filepath)
            )
        """)
        
        # Auto-migrate: Add description and is_readonly columns if they don't exist
        cursor = conn.cursor()
        cursor.execute("PRAGMA table_info(profiles)")
        columns = [info[1] for info in cursor.fetchall()]
        if "description" not in columns:
            cursor.execute("ALTER TABLE profiles ADD COLUMN description TEXT DEFAULT ''")
        if "is_readonly" not in columns:
            cursor.execute("ALTER TABLE profiles ADD COLUMN is_readonly BOOLEAN NOT NULL DEFAULT 0")
        
        # Create default profiles if the DB is empty
        cursor.execute("SELECT count(*) FROM profiles")
        if cursor.fetchone()[0] == 0:
            cursor.execute("INSERT INTO profiles (name, description, is_active, is_readonly) VALUES ('Vanilla (No Overlays)', 'Baseline ArduPilot build without any modifications. Read-only.', 1, 1)")
            cursor.execute("INSERT INTO profiles (name, description, is_active, is_readonly) VALUES ('Custom Build Alpha', 'Experimental tuning parameters.', 0, 0)")
        else:
            # Lock the existing Vanilla profile for users upgrading
            cursor.execute("UPDATE profiles SET is_readonly = 1 WHERE name = 'Vanilla (No Overlays)'")
            
        conn.commit()

init_db()

def get_db_connection():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

@app.get("/", response_class=HTMLResponse)
async def index(request: Request, success_msg: str = None, error_msg: str = None):
    conn = get_db_connection()
    
    profiles = conn.execute("SELECT * FROM profiles ORDER BY name").fetchall()
    active_profile = conn.execute("SELECT * FROM profiles WHERE is_active = 1").fetchone()
    
    files = []
    dirs = set()
    
    if active_profile:
        files = conn.execute("SELECT * FROM files WHERE profile_id = ? ORDER BY filepath", (active_profile['id'],)).fetchall()
        for f in files:
            dirname = os.path.dirname(f['filepath'])
            if dirname:
                dirs.add(dirname)
                
    conn.close()
                
    return templates.TemplateResponse("index.html", {
        "request": request, 
        "profiles": profiles,
        "active_profile": active_profile,
        "files": files,
        "dirs": sorted(list(dirs)),
        "success_msg": success_msg,
        "error_msg": error_msg
    })

@app.post("/profile/create")
async def create_profile(name: str = Form(...)):
    try:
        with get_db_connection() as conn:
            conn.execute("INSERT INTO profiles (name, description, is_active, is_readonly) VALUES (?, '', 0, 0)", (name.strip(),))
            conn.commit()
        msg = quote(f"Profile '{name}' created.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except sqlite3.IntegrityError:
        msg = quote("A profile with that name already exists.")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Profile creation failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to create profile.", status_code=303)

@app.post("/profile/rename")
async def rename_profile(profile_id: int = Form(...), new_name: str = Form(...)):
    try:
        with get_db_connection() as conn:
            prof = conn.execute("SELECT is_readonly FROM profiles WHERE id = ?", (profile_id,)).fetchone()
            if prof and prof['is_readonly']:
                return RedirectResponse(url=f"/patch-manager/?error_msg=Cannot rename a read-only profile.", status_code=303)
                
            conn.execute("UPDATE profiles SET name = ? WHERE id = ?", (new_name.strip(), profile_id))
            conn.commit()
        msg = quote(f"Profile renamed to '{new_name}'.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except sqlite3.IntegrityError:
        msg = quote("A profile with that name already exists.")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Profile rename failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to rename profile.", status_code=303)

@app.post("/profile/annotate")
async def annotate_profile(profile_id: int = Form(...), description: str = Form("")):
    try:
        with get_db_connection() as conn:
            prof = conn.execute("SELECT is_readonly FROM profiles WHERE id = ?", (profile_id,)).fetchone()
            if prof and prof['is_readonly']:
                return RedirectResponse(url=f"/patch-manager/?error_msg=Cannot edit notes on a read-only profile.", status_code=303)

            conn.execute("UPDATE profiles SET description = ? WHERE id = ?", (description, profile_id))
            conn.commit()
        msg = quote("Notes saved successfully.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Annotation failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to save notes.", status_code=303)

@app.post("/profile/duplicate")
async def duplicate_profile(profile_id: int = Form(...)):
    try:
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            src = cursor.execute("SELECT name, description FROM profiles WHERE id = ?", (profile_id,)).fetchone()
            if not src:
                raise Exception("Source profile not found.")
            
            base_name = src['name']
            desc = src['description'] if src['description'] else ""
            
            new_name = f"{base_name} (Copy)"
            counter = 1
            while True:
                exists = cursor.execute("SELECT id FROM profiles WHERE name = ?", (new_name,)).fetchone()
                if not exists:
                    break
                counter += 1
                new_name = f"{base_name} (Copy {counter})"
            
            cursor.execute("UPDATE profiles SET is_active = 0")
            
            # Duplicates are NEVER read-only, even if cloned from a read-only profile
            cursor.execute("INSERT INTO profiles (name, description, is_active, is_readonly) VALUES (?, ?, 1, 0)", (new_name, desc))
            new_profile_id = cursor.lastrowid
            
            files = cursor.execute("SELECT filepath, content FROM files WHERE profile_id = ?", (profile_id,)).fetchall()
            for f in files:
                cursor.execute("INSERT INTO files (profile_id, filepath, content) VALUES (?, ?, ?)", 
                               (new_profile_id, f['filepath'], f['content']))
            conn.commit()
            
        msg = quote(f"Seamlessly cloned and activated '{new_name}'.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Profile duplication failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to duplicate profile.", status_code=303)

@app.post("/profile/activate")
async def activate_profile(profile_id: int = Form(...)):
    try:
        with get_db_connection() as conn:
            conn.execute("UPDATE profiles SET is_active = 0")
            conn.execute("UPDATE profiles SET is_active = 1 WHERE id = ?", (profile_id,))
            conn.commit()
        return RedirectResponse(url=f"/patch-manager/", status_code=303)
    except Exception as e:
        logger.error(f"Profile activation failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to activate profile.", status_code=303)

@app.post("/profile/delete")
async def delete_profile(profile_id: int = Form(...)):
    try:
        with get_db_connection() as conn:
            prof = conn.execute("SELECT is_readonly, is_active FROM profiles WHERE id = ?").fetchone()
            
            if prof and prof['is_readonly']:
                msg = quote("Cannot delete a read-only profile.")
                return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)
                
            if prof and prof['is_active']:
                msg = quote("Cannot delete the currently active profile. Switch to another first.")
                return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)
                
            conn.execute("DELETE FROM profiles WHERE id = ?", (profile_id,))
            conn.commit()
        msg = quote("Profile deleted successfully.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Profile deletion failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to delete profile.", status_code=303)

@app.post("/deploy")
async def deploy_active_profile():
    try:
        with get_db_connection() as conn:
            active = conn.execute("SELECT id, name FROM profiles WHERE is_active = 1").fetchone()
            if not active:
                raise Exception("No active profile selected.")
                
            files = conn.execute("SELECT filepath, content FROM files WHERE profile_id = ?", (active['id'],)).fetchall()
            
        if os.path.exists(OVERLAY_DIR):
            for item in os.listdir(OVERLAY_DIR):
                item_path = os.path.join(OVERLAY_DIR, item)
                if os.path.isfile(item_path):
                    os.remove(item_path)
                elif os.path.isdir(item_path):
                    shutil.rmtree(item_path)
        os.makedirs(OVERLAY_DIR, exist_ok=True)
        
        count = 0
        for f in files:
            target = os.path.abspath(os.path.join(OVERLAY_DIR, f['filepath'].lstrip("/")))
            if target.startswith(OVERLAY_DIR):
                os.makedirs(os.path.dirname(target), exist_ok=True)
                with open(target, "wb") as out_file:
                    out_file.write(f['content'])
                count += 1
                
        return RedirectResponse(url="/", status_code=303)
    except Exception as e:
        logger.error(f"Deployment failed: {e}")
        msg = quote("Failed to deploy files to staging directory.")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

@app.post("/file/upload")
async def upload_file(
    file: UploadFile = File(...), 
    existing_path: str = Form(""), 
    new_path: str = Form(""), 
    profile_id: int = Form(...)
):
    try:
        with get_db_connection() as conn:
            prof = conn.execute("SELECT is_readonly FROM profiles WHERE id = ?", (profile_id,)).fetchone()
            if prof and prof['is_readonly']:
                msg = quote("Cannot upload files to a read-only profile.")
                return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

            target_path = new_path.strip() if new_path.strip() else existing_path.strip()
            clean_path = target_path.strip("/")
            
            safe_filename = os.path.basename(file.filename)
            full_filepath = os.path.join(clean_path, safe_filename) if clean_path else safe_filename
            
            content = await file.read()

            conn.execute("""
                INSERT INTO files (profile_id, filepath, content) 
                VALUES (?, ?, ?)
                ON CONFLICT(profile_id, filepath) 
                DO UPDATE SET content=excluded.content
            """, (profile_id, full_filepath, content))
            conn.commit()
            
        msg = quote(f"Successfully saved {safe_filename} to profile.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Upload failed: {e}")
        msg = quote("Failed to upload file to database.")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

@app.post("/file/delete")
async def delete_file(file_id: int = Form(...)):
    try:
        with get_db_connection() as conn:
            conn.execute("DELETE FROM files WHERE id = ?", (file_id,))
            conn.commit()
        return RedirectResponse(url=f"/patch-manager/", status_code=303)
    except Exception as e:
        logger.error(f"File deletion failed: {e}")
        return RedirectResponse(url=f"/patch-manager/?error_msg=Failed to delete file.", status_code=303)

@app.get("/edit", response_class=HTMLResponse)
async def edit_file(request: Request, file_id: int, error_msg: str = None):
    try:
        with get_db_connection() as conn:
            file_record = conn.execute("SELECT * FROM files WHERE id = ?", (file_id,)).fetchone()
            active_profile = conn.execute("SELECT * FROM profiles WHERE is_active = 1").fetchone()
            profiles = conn.execute("SELECT * FROM profiles ORDER BY name").fetchall()
            
        if not file_record:
            msg = quote("File not found in database.")
            return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

        try:
            content_text = file_record['content'].decode('utf-8')
        except UnicodeDecodeError:
            msg = quote(f"Safety Lock: {os.path.basename(file_record['filepath'])} appears to be a binary file and cannot be opened in the text editor.")
            return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

        return templates.TemplateResponse("edit.html", {
            "request": request, 
            "file_id": file_id,
            "filepath": file_record['filepath'], 
            "content": content_text,
            "active_profile": active_profile,
            "profiles": profiles,
            "error_msg": error_msg
        })
    except Exception as e:
        logger.error(f"Edit GET failed: {e}")
        msg = quote(f"System Error loading file: {str(e)}")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)

@app.post("/edit")
async def save_file(file_id: int = Form(...), content: str = Form(...)):
    try:
        content_bytes = content.encode('utf-8')
        with get_db_connection() as conn:
            conn.execute("UPDATE files SET content = ? WHERE id = ?", (content_bytes, file_id))
            conn.commit()
            
        msg = quote("Successfully saved changes.")
        return RedirectResponse(url=f"/patch-manager/?success_msg={msg}", status_code=303)
    except Exception as e:
        logger.error(f"Edit POST failed: {e}")
        msg = quote("Failed to save file.")
        return RedirectResponse(url=f"/patch-manager/edit?file_id={file_id}&error_msg={msg}", status_code=303)

@app.get("/download")
async def download_single_file(file_id: int):
    try:
        with get_db_connection() as conn:
            file_record = conn.execute("SELECT * FROM files WHERE id = ?", (file_id,)).fetchone()
            
        if not file_record:
            raise HTTPException(status_code=404, detail="File not found")
            
        with tempfile.NamedTemporaryFile(delete=False) as tmp:
            tmp.write(file_record['content'])
            tmp_path = tmp.name
            
        return FileResponse(
            path=tmp_path, 
            filename=os.path.basename(file_record['filepath']),
            media_type='application/octet-stream'
        )
    except Exception as e:
        logger.error(f"Download single file failed: {e}")
        msg = quote("Failed to download file.")
        return RedirectResponse(url=f"/patch-manager/?error_msg={msg}", status_code=303)