refactor: перестроена файловая структура в виде python-пакета с файлом main.py
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user