304 lines
8.7 KiB
Python
304 lines
8.7 KiB
Python
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()
|