diff --git a/src/traffic_checker.py b/src/traffic_checker.py deleted file mode 100644 index d33275e..0000000 --- a/src/traffic_checker.py +++ /dev/null @@ -1,51 +0,0 @@ -import json -import os -from pathlib import Path - -import requests - - -def load_config() -> dict: - dev_config = Path(__file__).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 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() - - -if __name__ == "__main__": - config = load_config() - - for server in config["servers"]: - data = fetch_server_data(server["uuid"]) - print(data) diff --git a/src/traffic_checker/__init__.py b/src/traffic_checker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/traffic_checker/main.py b/src/traffic_checker/main.py new file mode 100644 index 0000000..4d62a55 --- /dev/null +++ b/src/traffic_checker/main.py @@ -0,0 +1,283 @@ +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 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, + ) + + state[server["uuid"]] = event["server_state"] + + save_state(state_file_path, 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) + + print(f"""Server {server["target"]}: {event["event"]}""") + + except Exception as e: + print(f"""Error processing server {server["target"]}: {e}""") + + +if __name__ == "__main__": + main()