Spaces:

omnitool-ai
/

omnitool_on_hf

Running

File size: 11,396 Bytes

const http = require('http');
const { spawn } = require('child_process');

const PROXY_PORT_OMNITOOL = 4444;
const CONTAINER_HOST = '127.0.0.1';
const CONTAINER_PORT_OMNITOOL = 1688;
const PROTOCOL = 'https';
let OMNITOOL_READY = false;
let OMNITOOL_PRE_READY = false;
let ALREADY_STARTING = false;
let IFRAME = false;
const OMNITOOL_STATUS_CHECK_INTERVAL = 2000; // 2 seconds

const VERSION = '0.0.7a';
console.log(`************ Omnitool Proxy Server v${VERSION} ************`);
let omnitoolLogs = [];

// HTML templates and constants
const START_SERVER_BUTTON_HTML = `<button id="startServerButton" onclick="startServer()">Start Omnitool Server</button>`;
const GOTO_OMNITOOL_BUTTON_HTML = `<button id="exitIframeButton" onclick="exitIframe()">GOTO OMNITOOL</button>`;
const LOG_CONTAINER_HTML = `<div id="logContainer"></div>`;
const COMMON_STYLES = `
    <style>
        #logContainer {
            height: 400px;
            overflow-y: scroll;
            background-color: black;
            color: lime;
            font-family: 'Courier New', Courier, monospace;
            padding: 10px;
            white-space: pre-wrap;
            border: 1px solid #ddd;
        }
        .highlight-button {
            animation: pulseAnimation 1s infinite;
            background-color: yellow;
            color: black;
            font-weight: bold;
        }                    
        @keyframes pulseAnimation {
            0% { transform: scale(1); }
            50% { transform: scale(1.1); }
            100% { transform: scale(1); }
        }
    </style>`;

// Function to generate HTML content based on server status
function generateHTMLContent(newUrl) {
    let startButtonDisabled = OMNITOOL_PRE_READY || OMNITOOL_READY || ALREADY_STARTING ? 'disabled' : '';
    let startButtonClass = !startButtonDisabled ? 'highlight-button' : '';
    let gotoButtonClass = IFRAME ? 'highlight-button' : '';

    // set the startButton to disabled if OMNITOOL_PRE_READY or OMNITOOL_READY or ALREADY_STARTING are true
    let buttonsHTML = `<button id="startServerButton" class="${startButtonClass}" onclick="startServer()" ${startButtonDisabled}>Start Omnitool Server</button>`;

    if (IFRAME) 
    {
        buttonsHTML += `<button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>`;
    }

    return `
        <html>
            <head>
                <title>Proxy Server</title>
                ${COMMON_STYLES}
            </head>
            <body>
                ${buttonsHTML}
                ${LOG_CONTAINER_HTML}
                <script>
                    function startServer() {
                        document.getElementById('startServerButton').classList.remove
                        ('highlight-button');
                        document.getElementById('startServerButton').disabled = true;


                        fetch('/start-omnitool-server')
                            .then(response => response.json())
                            .then(data => startLogPolling());
  
    
                    }

                    function exitIframe() {
                        window.open('${newUrl}', '_blank');
                    }

                    function startLogPolling() {

                        const interval = setInterval(() => {
                            fetch('/omnitool-logs')
                                .then(response => response.json())
                                .then(data => {
                                    const logContainer = document.getElementById('logContainer');
                                    logContainer.innerText = data.logs.join("\\n");
                                    if (data.ready) {
                                        clearInterval(interval);
                                    }
                                    else {
                                        scrollToBottom(logContainer);
                                    }
                                });
                        }, 500);
                    }

                    function scrollToBottom(element) {
                        element.scrollTop = element.scrollHeight;
                    }

                    startLogPolling();
                </script>
            </body>
        </html>`;
}


async function startOmnitoolServer()
{
    if (ALREADY_STARTING) return;
    ALREADY_STARTING = true;

    console.log('Starting Omnitool Server...');
    return new Promise((resolve, reject) =>
    {
        const omnitoolStartProcess = spawn('./omnitool_start.sh');

        omnitoolStartProcess.stdout.on('data', (data) =>
        {
            omnitoolLogs.push(data.toString());
            console.log(`omnitool stdout: ${data}`);
        });

        omnitoolStartProcess.stderr.on('data', (data) =>
        {
            console.error(`omnitool stderr: ${data}`);
            omnitoolLogs.push(data.toString());
        });

        omnitoolStartProcess.on('close', (code) =>
        {
            const message = `Omnitool server process exited with code: ${code}`;
            console.log(message);
            omnitoolLogs.push(message);

            ALREADY_STARTING = false;
            if (code === 0)
            {
                resolve();
            } else
            {
                reject(`Omnitool server did not start properly.`);
            }
        });
    });
}

function setOmnitoolReady()
{
    if (OMNITOOL_PRE_READY === false) 
    {
        OMNITOOL_PRE_READY = true;
        console.log('Omnitool server is (PRE) ready! ');
        setTimeout(() =>
        {
            console.log('Omnitool server is (NOW) ready! ');
            OMNITOOL_READY = true;
        }, 2000); // Delay by 2 second
    }
}

async function periodicCheckOmnitoolStatus() {
    console.log("Starting async operation...");
    try
    {
        await checkOmnitoolStatus();
    }
    catch (error)
    {
        console.log(`Omnitool not ready: ${error} (IFRAME = ${IFRAME}`);
        OMNITOOL_READY = false;
        OMNITOOL_PRE_READY = false;
    }
    // Schedule the next execution
    setTimeout(periodicCheckOmnitoolStatus, OMNITOOL_STATUS_CHECK_INTERVAL);
}


async function checkOmnitoolStatus()
{

    console.log(`BEFORE: Checking Omnitool status: PRE_READY = ${OMNITOOL_PRE_READY}, READY = ${OMNITOOL_READY}, IFRAME = ${IFRAME}, ALREADY_STARTING = ${ALREADY_STARTING}`);

    return new Promise((resolve, reject) =>
    {
        const options = {
            hostname: CONTAINER_HOST,
            port: CONTAINER_PORT_OMNITOOL,
            path: '/',
            method: 'GET'
        };

        const req = http.request(options, (res) =>
        {
            if (res.statusCode === 200)
            {
                IFRAME = false;
                setOmnitoolReady();
            } 
            else
            {
                OMNITOOL_READY = false;
                OMNITOOL_PRE_READY = false;
            }
            resolve();
        });

        req.on('error', (error) =>
        {
            if (error.code === 'ECONNREFUSED')
            {
                IFRAME  = true;
            }
        });

        req.end();
    });


}



async function startRequestForwardingServer(proxy_port)
{
    const server = http.createServer(async (req, res) =>
    {

        console.log("------------------");
        const localUrl = req.headers['host'];
        const hostname = localUrl.split(':')[0];
        const newUrl = `${PROTOCOL}://${hostname}:${PROXY_PORT_OMNITOOL}`;

        if (!OMNITOOL_READY || IFRAME)
        {
            // Handle GET requests here
            console.log(`Omnitool Server:OMNITOOL_READY =  ${OMNITOOL_READY}`);
            if (req.method === 'GET')
            {
                switch (req.url)
                {
                    case '/start-omnitool-server':
                        {

                            console.log(`Omnitool Server:ALREADY_STARTING =  ${ALREADY_STARTING}`);
                            if (ALREADY_STARTING) 
                            {
                                res.writeHead(200, { 'Content-Type': 'text/html' });
                                res.end("Omnitool server already starting");
                                return;
                            }
                            try
                            {
                                await startOmnitoolServer();
                                res.writeHead(200, { 'Content-Type': 'text/html' });
                                res.end("Omnitool server started successfully");
                            }
                            catch (error)
                            {
                                console.error(error);
                                ALREADY_STARTING = false;
                                res.writeHead(500, { 'Content-Type': 'text/html' });
                                res.end(`Error starting Omnitool server: ${error}`);
                            }

                            return;
                        }

                    case '/omnitool-logs':
                        {
                            res.writeHead(200, { 'Content-Type': 'application/json' });
                            const reply = { logs: omnitoolLogs, ready: OMNITOOL_PRE_READY };
                            res.end(JSON.stringify(reply));
                            return;
                        }

                    case '/':
                        {
                            res.writeHead(200, { 'Content-Type': 'text/html' });
                            res.end(generateHTMLContent(newUrl));
                            return;
                        }

                    default:
                        res.writeHead(404, { 'Content-Type': 'text/html' });
                        res.end(`<html><body><h1>404 Not Found</h1><p>The requested URL ${req.url} was not found on this server.</p></body></html>`);
                        return;
                }
            }
            else 
            {
                // Handle non-GET requests here
                res.writeHead(405, { 'Content-Type': 'text/html' });
                res.end(`<html><body><h1>405 Method Not Allowed</h1><p>The request method ${req.method} is not allowed for the URL ${req.url}.</p></body></html>`);
            }
        }
        else
        {
            // Proxy logic...
            const options = { hostname: CONTAINER_HOST, port: CONTAINER_PORT_OMNITOOL, path: req.url, method: req.method, headers: req.headers };
            const proxy = http.request(options, (proxyRes) =>
            {
                res.writeHead(proxyRes.statusCode, proxyRes.headers);
                proxyRes.pipe(res, { end: true });
            });
            req.pipe(proxy, { end: true });
        }
    });

    server.listen(proxy_port, '0.0.0.0');
    console.log(`Request forwarding server listening on port ${proxy_port}`);
}

async function startManagementServer()
{
    try
    {
        await startRequestForwardingServer(PROXY_PORT_OMNITOOL);

        // Set the interval in milliseconds (e.g., 2000 milliseconds = 2 seconds)

        // Use setInterval to schedule the function
        setInterval(checkOmnitoolStatus, OMNITOOL_STATUS_CHECK_INTERVAL);
    } catch (error)
    {
        console.error('Failed to start servers:', error);
    }
}

startManagementServer();