import json import os from datetime import datetime, timedelta from pathlib import Path from zoneinfo import ZoneInfo import requests def load_config() -> dict: dev_config = Path(__file__).parent.parent.parent / "config.json" sys_config = Path("/etc/traffic-checker/config.json") if dev_config.exists(): config_path = dev_config elif sys_config.exists(): config_path = sys_config else: raise FileNotFoundError( "Configuration file config.json not found in dev path or /etc/traffic-checker/" ) with open(config_path, "r", encoding="utf-8") as f: return json.load(f) def get_state_file_path() -> Path: dev_state = Path(__file__).parent.parent.parent / "state.json" sys_state = Path("/var/lib/traffic-checker/state.json") if dev_state.exists() or not sys_state.parent.exists(): return dev_state return sys_state def load_state(path: Path) -> dict: if not path.exists(): return {} try: with open(path, "r", encoding="utf-8") as f: return json.load(f) except json.JSONDecodeError: return {} def save_state(path: Path, state: dict) -> None: with open(path, "w", encoding="utf-8") as f: json.dump(state, f, indent=2, ensure_ascii=False) def fetch_server_data(server_uuid: str) -> dict: token = os.getenv("VIRTFUSION_TOKEN") if not token: raise RuntimeError("VIRTFUSION_TOKEN environment variable is not set.") url = f"https://virtfusion-nat.gullo.me/api/server/{server_uuid}" headers = { "Accept": "application/json", "Authorization": f"Bearer {token}", } params = {"state": "true"} response = requests.get(url, headers=headers, params=params, timeout=10) response.raise_for_status() return response.json() def process_traffic_stats(server_data: dict) -> dict: traffic_limit_str = str(server_data["data"]["network"]["primary"]["limit"]) traffic_limit_gbyte = float(traffic_limit_str.removesuffix("GB").strip()) traffic_usage_byte = int( server_data["data"]["state"]["network"]["primary"]["traffic"]["total"] ) traffic_usage_gbyte = traffic_usage_byte / 1024**3 traffic_usage_percent = traffic_usage_gbyte / traffic_limit_gbyte * 100 return { "limit": traffic_limit_gbyte, "used": round(traffic_usage_gbyte, 2), "percentage": round(traffic_usage_percent, 2), } def process_monthly_period( server_data: dict, timezone: str = "Europe/Moscow" ) -> dict: period_start_utc = datetime.fromisoformat( server_data["data"]["currentMonthlyPeriod"]["start"] ) period_end_utc = datetime.fromisoformat( server_data["data"]["currentMonthlyPeriod"]["end"] ) period_reset_utc = period_end_utc + timedelta(microseconds=1) return { "start": period_start_utc.astimezone(ZoneInfo(timezone)), "end": period_end_utc.astimezone(ZoneInfo(timezone)), "reset": period_reset_utc.astimezone(ZoneInfo(timezone)), } def process_traffic_event( server_uuid: str, warning_thresholds: list[float], traffic_percentage: float, period_reset_iso: str, state: dict, ) -> dict: server_state = state.get( server_uuid, { "last_threshold": 0.0, "is_critical": False, "period_reset": period_reset_iso, }, ) if server_state.get("period_reset") != period_reset_iso: new_server_state = { "last_threshold": 0.0, "is_critical": False, "period_reset": period_reset_iso, } was_critical = server_state.get("is_critical", False) return { "event": "reset" if was_critical else "none", "server_state": new_server_state, } last_threshold = server_state.get("last_threshold", 0) if traffic_percentage >= 100.0: if last_threshold < 100.0: new_server_state = { "last_threshold": 100.0, "is_critical": True, "period_reset": period_reset_iso, } return { "event": "critical", "server_state": new_server_state, } passed_thresholds = [ t for t in warning_thresholds if traffic_percentage >= t ] if passed_thresholds: max_passed = max(passed_thresholds) if max_passed > last_threshold: new_server_state = { "last_threshold": max_passed, "is_critical": False, "period_reset": period_reset_iso, } return {"event": "warning", "server_state": new_server_state} return {"event": "none", "server_state": server_state} def format_message( event_type: str, server_target: str, server_pretty_name: str, traffic_percentage: float, period_reset: datetime, ) -> str: if event_type == "warning": return f"⚠️ На сервере {server_target} ({server_pretty_name}) израсходовано {traffic_percentage}% трафика." MONTHS_RU = { 1: "января", 2: "февраля", 3: "марта", 4: "апреля", 5: "мая", 6: "июня", 7: "июля", 8: "августа", 9: "сентября", 10: "октября", 11: "ноября", 12: "декабря", } period_reset_str = f"{period_reset.day} {MONTHS_RU[period_reset.month]} в {period_reset.strftime("%H:%M")}" if event_type == "critical": return ( f"⛔️ На сервере {server_target} ({server_pretty_name}) введено ограничение скорости в 625 килобит/секунду из-за исчерпания объема трафика.\n\n" f"Снятие ограничений скорости состоится {period_reset_str}." ) if event_type == "reset": return f"🔄 На сервере {server_target} ({server_pretty_name}) снято ограничение скорости." raise ValueError( f"""Unknown event type: "{event_type}". """ f"""Expected one of: "warning", "critical", "reset".""" ) def send_telegram_message( message: str, chat_id: int | str | None = None ) -> None: token = os.getenv("TELEGRAM_BOT_TOKEN") if not token: raise RuntimeError( "TELEGRAM_BOT_TOKEN environment variable is not set." ) target_chat_id = chat_id or os.getenv("TELEGRAM_CHAT_ID") if not target_chat_id: raise RuntimeError("TELEGRAM_CHAT_ID is not provided.") url = f"https://api.telegram.org/bot{token}/sendMessage" payload = { "chat_id": target_chat_id, "text": message, } response = requests.post(url, json=payload, timeout=10) response.raise_for_status() def to_log( server_target: str, event_type: str, traffic_stats: dict, monthly_period: dict, ): period_start_str = monthly_period["start"].strftime("%Y-%m-%d %H:%M:%S") period_reset_str = monthly_period["reset"].strftime("%Y-%m-%d %H:%M:%S") print(f"Traffic check for server: {server_target}") print(f"Event: {event_type}") print( f"Traffic stats: {traffic_stats["percentage"]}%, " f"{traffic_stats["used"]} GB / {traffic_stats["limit"]} GB" ) print(f"Monthly period: {period_start_str} -> {period_reset_str}") def main(): config = load_config() state_file_path = get_state_file_path() state = load_state(state_file_path) for server in config["servers"]: try: data = fetch_server_data(server["uuid"]) traffic_stats = process_traffic_stats(data) monthly_period = process_monthly_period(data) event = process_traffic_event( server["uuid"], config["thresholds"]["warning_percents"], traffic_stats["percentage"], monthly_period["reset"].isoformat(), state, ) if event["event"] != "none": message = format_message( event["event"], server["target"], server["pretty_name"], traffic_stats["percentage"], monthly_period["reset"], ) send_telegram_message(message) to_log( server["target"], event["event"], traffic_stats, monthly_period ) state[server["uuid"]] = event["server_state"] save_state(state_file_path, state) except Exception as e: print(f"""Error processing server {server["target"]}: {e}""") if __name__ == "__main__": main()