view collector/racc.py @ 4:5ad799e534a5

move to subdir
author drewp@bigasterisk.com
date Mon, 13 Mar 2023 18:50:23 -0700
parents racc.py@3d7fc94a404a
children
line wrap: on
line source

import logging
import os
import socket
import sys

import background_loop
import daemonocle
import psutil
import uvicorn
from prometheus_client import Gauge
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import HTMLResponse
from starlette.routing import Route
from starlette.staticfiles import StaticFiles
from starlette_exporter import PrometheusMiddleware, handle_metrics
''
import progs_all as progs

if psutil.OSX:
    import current_window_title_osx as current_window_title
    import idle_osx as idle
    import power_osx as power
    import volume_osx as volume
elif psutil.LINUX:
    import current_window_title_linux as current_window_title
    import idle_linux as idle
    import power_linux as power
    import volume_linux as volume
else:
    raise NotImplementedError(repr(sys.implementation))

DEBUG = os.environ.get('RACC_DEBUG', False)
logging.basicConfig(level=logging.INFO)
log = logging.getLogger()
hostname = socket.gethostname().split('.')[0]
RACC_RUNNING = Gauge("racc_running", "program is running", ['host', 'prog'])
RACC_IDLE = Gauge("racc_idle", "desktop mouse/kb idle seconds", ['host'])
RACC_SCREEN = Gauge("racc_screen_on", "screen in unlocked/on mode", ['host'])
RACC_CURRENT_WINDOW = Gauge("racc_current_window",
                            "label carries title; site is last part",
                            ['host', 'title', 'site'])


def update(first_run):
    try:
        for p, val in progs.get_running_progs().items():
            RACC_RUNNING.labels(host=hostname, prog=p).set(val)
    except Exception:
        if DEBUG: raise

    try:
        RACC_IDLE.labels(host=hostname).set(idle.get_idle_seconds())
    except Exception:
        if DEBUG: raise

    try:
        RACC_SCREEN.labels(host=hostname).set(power.is_screen_on())
    except Exception:
        if DEBUG: raise

    try:
        title = current_window_title.get_current_window_title()
        # chrome, at least on osx, adds icon to window title when it's playing audio
        title = title.rstrip('🔊').strip()
        # some websites choose to title like '<something> - <site name>'
        last_section = title.split(' - ')[-1].strip()
        if last_section in {'YouTube', 'Google Search', 'Google Chrome', 'Visual Studio Code', 'Roblox'}:
            site = last_section.lower().replace(' ', '_')
        else:
            site = ''
        RACC_CURRENT_WINDOW.clear()
        RACC_CURRENT_WINDOW.labels(host=hostname, title=title,
                                   site=site).set(1)
    except Exception:
        if DEBUG: raise


async def root(req: Request) -> HTMLResponse:
    vol = await volume.get_volume()
    return HTMLResponse(f'''controls for {hostname} whose volume is {vol}''')


def main():
    app = Starlette(debug=True,
                    routes=[
                        Route('/', root),
                    ],
                    on_startup=[
                        lambda: background_loop.loop_forever(update, 5),
                    ])

    app.add_middleware(PrometheusMiddleware, app_name='racc')
    app.add_route("/metrics", handle_metrics)
    uvicorn.run(app, host='0.0.0.0', port=5150)


if __name__ == "__main__":
    d = daemonocle.Daemon(
        worker=main,
        work_dir='.',
        pid_file='/tmp/racc.pid',
        detach=True,
    )
    d.stop(timeout=1, force=True)
    if DEBUG:
        main()
    else:
        d.cli()