10
0

Enforce Prometheus bearer token usage and switch to target GET parameter

This commit is contained in:
Lauri Võsandi 2021-07-04 10:58:52 +03:00
parent fafd1db041
commit 192f427b35

View File

@ -10,14 +10,13 @@ PREFIX = os.getenv("PROMETHEUS_PREFIX", "mikrotik_")
PROMETHEUS_BEARER_TOKEN = os.getenv("PROMETHEUS_BEARER_TOKEN") PROMETHEUS_BEARER_TOKEN = os.getenv("PROMETHEUS_BEARER_TOKEN")
MIKROTIK_USER = os.getenv("MIKROTIK_USER") MIKROTIK_USER = os.getenv("MIKROTIK_USER")
MIKROTIK_PASSWORD = os.getenv("MIKROTIK_PASSWORD") MIKROTIK_PASSWORD = os.getenv("MIKROTIK_PASSWORD")
TARGETS = os.getenv("TARGETS")
if not MIKROTIK_USER: if not MIKROTIK_USER:
raise ValueError("MIKROTIK_USER not specified") raise ValueError("MIKROTIK_USER not specified")
if not MIKROTIK_PASSWORD: if not MIKROTIK_PASSWORD:
raise ValueError("MIKROTIK_PASSWORD not specified") raise ValueError("MIKROTIK_PASSWORD not specified")
if not TARGETS: if not PROMETHEUS_BEARER_TOKEN:
raise ValueError("TARGETS not specified") raise ValueError("No PROMETHEUS_BEARER_TOKEN specified")
RATE_MAPPING = { RATE_MAPPING = {
"40Gbps": 40 * 10 ** 9, "40Gbps": 40 * 10 ** 9,
@ -165,14 +164,11 @@ async def scrape_mikrotik(target):
@app.route("/metrics") @app.route("/metrics")
async def view_export(request): async def view_export(request):
if PROMETHEUS_BEARER_TOKEN and request.token != PROMETHEUS_BEARER_TOKEN: if request.token != PROMETHEUS_BEARER_TOKEN:
raise exceptions.Forbidden("Invalid bearer token") raise exceptions.Forbidden("Invalid bearer token")
async def streaming_fn(response): async def streaming_fn(response):
args = [scrape_mikrotik(target) for target in TARGETS.split(",")] async for line in wrap(scrape_mikrotik(request.args.get("target"))):
combine = stream.merge(*args)
async with combine.stream() as streamer:
async for line in wrap(streamer):
await response.write(line + "\n") await response.write(line + "\n")
return response.stream(streaming_fn, content_type="text/plain") return response.stream(streaming_fn, content_type="text/plain")