You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
3.1 KiB
90 lines
3.1 KiB
from fastapi import FastAPI
|
|
from fastapi import Request
|
|
from fastapi.responses import Response
|
|
from fastapi import BackgroundTasks
|
|
import uvicorn
|
|
import os
|
|
from time import time
|
|
import asyncio
|
|
import datetime
|
|
|
|
class Collector:
|
|
app = FastAPI()
|
|
secrets = []
|
|
hosts = []
|
|
store = {}
|
|
def __init__(self) -> None:
|
|
self.secrets = os.getenv("KEYS", "secret").split(",")
|
|
self.hosts = os.getenv("HOSTS", "srv").split(",")
|
|
self.dead_after = int(os.getenv("DEAD_AFTER", "90"))
|
|
self.build_routes()
|
|
|
|
def check_headers(self, request: Request):
|
|
if not request.headers.get("Secret", ""):
|
|
print("secret empty")
|
|
return Response(status_code=400)
|
|
if not request.headers.get("Host", "") and not request.headers.get("Host-replace", ""):
|
|
print("host empty")
|
|
return Response(status_code=400)
|
|
|
|
if not request.headers.get("Secret", "") in self.secrets:
|
|
return Response(status_code=400)
|
|
|
|
host = request.headers.get("Host-replace", "") or request.headers.get("Host", "")
|
|
host = host.split(":")[0]
|
|
if not host in self.hosts:
|
|
return Response(status_code=400)
|
|
|
|
return None
|
|
|
|
async def deadline(self, host):
|
|
await asyncio.sleep(self.dead_after)
|
|
if host in self.store:
|
|
if time()+1 - self.store[host]['timestamp'] > self.dead_after:
|
|
print(f"{host} is dead")
|
|
del self.store[host]
|
|
|
|
|
|
def build_routes(self):
|
|
@self.app.get("/metrics")
|
|
async def get(request: Request):
|
|
response: Response = self.check_headers(request)
|
|
if response:
|
|
return response
|
|
|
|
host = request.headers.get("Host-replace", "") or request.headers.get("Host", "")
|
|
host = host.split(":")[0]
|
|
|
|
if host in self.store and self.store[host]['data']:
|
|
return Response(
|
|
content=self.store[host]['data'],
|
|
status_code=200,
|
|
headers=self.store[host]['headers'])
|
|
else:
|
|
print(f"cannot find {host}")
|
|
return Response(status_code=404)
|
|
|
|
@self.app.post("/metrics")
|
|
async def post(request: Request, background_tasks: BackgroundTasks):
|
|
response: Response = self.check_headers(request)
|
|
if response:
|
|
return response
|
|
|
|
host = request.headers.get("Host-replace", "") or request.headers.get("Host", "")
|
|
host = host.split(":")[0]
|
|
|
|
self.store[host] = {
|
|
"data":await request.body(),
|
|
"timestamp": time(),
|
|
"headers": {
|
|
"Content-Type": "text/plain; version=0.0.4; charset=utf-8; escaping=underscores",
|
|
"Date": datetime.datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S GMT")
|
|
}
|
|
}
|
|
|
|
background_tasks.add_task(self.deadline, host)
|
|
return Response(status_code=200)
|
|
|
|
if __name__ == "__main__":
|
|
collector = Collector()
|
|
uvicorn.run(collector.app, host = os.getenv("NEC_HOST", "0.0.0.0"), port = int(os.getenv("NEC_PORT", "9100")))
|