app.py

import os
import gradio as gr
import threading
from flask import Flask, request
from psycopg2 import connect, sql
import psycopg2.extras

# Fetch environment variables
DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")
APP_PASSWORD = os.getenv("APP_PASSWORD")

# Database connection function
def get_db_connection():
    return connect(
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD,
        host=DB_HOST,
        port=DB_PORT
    )

# Dictionary to manage server threads
server_threads = {}

# Flask app template for HTTP servers
def create_server(name, port):
    app = Flask(name)

    @app.route("/")
    def home():
        return f"Server {name} is running on port {port}!"

    return app

# Function to start a server
def start_server(name, port):
    if name in server_threads:
        return f"Server {name} is already running!"
    
    app = create_server(name, port)

    def run():
        app.run(host="0.0.0.0", port=port, debug=False, use_reloader=False)

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    server_threads[name] = thread

    # Update database status
    conn = get_db_connection()
    with conn.cursor() as cursor:
        cursor.execute(
            "INSERT INTO servers (name, port, status) VALUES (%s, %s, %s) ON CONFLICT (name) DO UPDATE SET status = EXCLUDED.status",
            (name, port, "running")
        )
        conn.commit()
    conn.close()

    return f"Server {name} started on port {port}!"

# Function to stop a server (dummy implementation for now)
def stop_server(name):
    if name not in server_threads:
        return f"Server {name} is not running!"

    # Terminate the thread (not directly supported, requires external handling)
    server_threads.pop(name)
    return f"Server {name} stopped!"

# Gradio UI for managing servers
def manage_servers(password, action, name, port):
    if password != APP_PASSWORD:
        return "Invalid password!"
    if action == "start":
        return start_server(name, port)
    elif action == "stop":
        return stop_server(name)
    else:
        return "Invalid action!"

# Function to handle HTTP commands
def execute_db_command(password, command):
    if password != APP_PASSWORD:
        return "Invalid password!"

    # Establish DB connection
    conn = get_db_connection()
    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
    
    try:
        # Execute the command (SQL injection protection by using parameterized queries)
        cursor.execute(command)
        result = cursor.fetchall()
        conn.commit()
    except Exception as e:
        result = f"Error executing command: {e}"
    finally:
        conn.close()
    
    return result

# Gradio UI for managing servers and interacting with DB
def get_db_stats():
    conn = get_db_connection()
    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
    
    # CockroachDB query to get database size (total used space in bytes)
    cursor.execute("""
    SELECT 
        db_name,
        SUM(range_size) AS total_size_bytes
    FROM crdb_internal.database_statistics
    WHERE db_name = current_database()
    GROUP BY db_name;
    """)
    db_stats = cursor.fetchone()

    conn.close()
    if db_stats:
        return f"Database {db_stats['db_name']} size: {db_stats['total_size_bytes']} bytes"
    else:
        return "Unable to fetch database size"

# Gradio Interface
with gr.Blocks() as server_manager:
    gr.Markdown("# ZeroCPU Server Manager")

    # Gradio UI for managing servers
    with gr.Row():
        password = gr.Textbox(label="Password", type="password")
    with gr.Row():
        action = gr.Radio(["start", "stop"], label="Action")
        name = gr.Textbox(label="Server Name")
        port = gr.Number(label="Port")
    with gr.Row():
        submit = gr.Button("Submit")
    output = gr.Textbox(label="Output")

    # Submit action for server management
    submit.click(
        manage_servers,
        inputs=[password, action, name, port],
        outputs=output
    )

    # Gradio UI for viewing DB stats and executing commands
    with gr.Row():
        command_input = gr.Textbox(label="Database Command", placeholder="Enter SQL command")
        command_submit = gr.Button("Execute Command")
    command_output = gr.Textbox(label="Command Output")

    # Submit action for executing DB commands
    command_submit.click(
        execute_db_command,
        inputs=[password, command_input],
        outputs=command_output
    )

    # Show the database size
    with gr.Row():
        db_stats_output = gr.Textbox(label="Database Stats")
    gr.Button("Show Database Stats").click(
        get_db_stats,
        outputs=db_stats_output
    )

    # Display message when Gradio app is running
    gr.Markdown("ZeroCPU is active")

server_manager.launch()