forked from zovos/vk_hackathon
107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
TCP log server — receives JSON-line logs from index-service and search-service.
|
|
|
|
Usage:
|
|
python3 server.py # listen on 0.0.0.0:9999
|
|
python3 server.py --port 9999
|
|
python3 server.py --save logs.jsonl # also save to file
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import socketserver
|
|
import sys
|
|
import threading
|
|
from datetime import datetime
|
|
|
|
COLORS = {
|
|
"DEBUG": "\033[36m",
|
|
"INFO": "\033[0m",
|
|
"WARNING": "\033[33m",
|
|
"ERROR": "\033[31m",
|
|
"CRITICAL": "\033[35m",
|
|
}
|
|
RESET = "\033[0m"
|
|
SERVICE_COLOR = {
|
|
"index-service": "\033[34m", # blue
|
|
"search-service": "\033[32m", # green
|
|
}
|
|
|
|
_save_file = None
|
|
_save_lock = threading.Lock()
|
|
|
|
|
|
def _format(entry: dict) -> str:
|
|
ts = entry.get("ts", "")[:23].replace("T", " ")
|
|
level = entry.get("level", "INFO")
|
|
service = entry.get("service", "?")
|
|
msg = entry.get("msg", "")
|
|
|
|
lc = COLORS.get(level, "")
|
|
sc = SERVICE_COLOR.get(service, "\033[0m")
|
|
return f"{ts} {sc}{service:<15}{RESET} {lc}{level:<8}{RESET} {msg}"
|
|
|
|
|
|
def _handle_line(raw: str) -> None:
|
|
raw = raw.strip()
|
|
if not raw:
|
|
return
|
|
try:
|
|
entry = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
entry = {"ts": datetime.utcnow().isoformat(), "level": "INFO", "service": "?", "msg": raw}
|
|
|
|
print(_format(entry), flush=True)
|
|
|
|
if _save_file:
|
|
with _save_lock:
|
|
_save_file.write(raw + "\n")
|
|
_save_file.flush()
|
|
|
|
|
|
class _Handler(socketserver.StreamRequestHandler):
|
|
def handle(self) -> None:
|
|
addr = self.client_address[0]
|
|
print(f"\033[90m[+] connected: {addr}{RESET}", flush=True)
|
|
try:
|
|
for raw_bytes in self.rfile:
|
|
try:
|
|
_handle_line(raw_bytes.decode("utf-8", errors="replace"))
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
print(f"\033[90m[-] disconnected: {addr}{RESET}", flush=True)
|
|
|
|
|
|
def main() -> None:
|
|
global _save_file
|
|
|
|
parser = argparse.ArgumentParser(description="TCP JSON-line log receiver")
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
parser.add_argument("--port", type=int, default=9999)
|
|
parser.add_argument("--save", metavar="FILE", help="Also save raw JSON lines to this file")
|
|
args = parser.parse_args()
|
|
|
|
if args.save:
|
|
_save_file = open(args.save, "a", encoding="utf-8")
|
|
print(f"Saving logs to {args.save}", flush=True)
|
|
|
|
server = socketserver.ThreadingTCPServer((args.host, args.port), _Handler)
|
|
server.allow_reuse_address = True
|
|
|
|
print(f"Listening on {args.host}:{args.port} ...\n", flush=True)
|
|
try:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
print("\nStopped.")
|
|
finally:
|
|
server.server_close()
|
|
if _save_file:
|
|
_save_file.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|