667 lines
28 KiB
Python
667 lines
28 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Royirin Rosmontis Private Proxy JP 流量监控（Azure 门户口径版）
|
||
- 主计量口径：Azure 门户免费服务页 "Data Transfer Out (GB)"（Standard Data Transfer Out - Free，MGN 路由）
|
||
—— 通过 VM 托管身份调用 Consumption Usage Details API 拉取，按自然月聚合，缓存至 data/azure_usage.json
|
||
- 参考口径：vnstat -i eth0（实时原始出站，仅作参考展示）
|
||
- 数据源（参考）：vnstat -i <interface> --json m/d（eth0 为 bond 主接口，覆盖 nginx 443 与 sing-box 8443 全部流量）
|
||
- 按 年/月 聚合成 JSON 存入 data/YYYY/YYYY-MM.json（弃用 SQLite）
|
||
- HTTP 服务: GET /vpn/status/traffic 返回当月用量 JSON（Azure 口径为主）
|
||
- 阈值告警: 达到配额 50/80/90/95/99/100% 时发邮件（每档每月只发一次，跳档补发；基于 Azure 口径）
|
||
- 月报: 每月 1 日 00:05 起发送上个月报告（失败重试；基于 Azure 口径）
|
||
"""
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import socket
|
||
import smtplib
|
||
import ssl
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
import urllib.parse
|
||
import urllib.request
|
||
from datetime import date, datetime, timedelta, timezone
|
||
from email.message import EmailMessage
|
||
from email.utils import formataddr
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
log = logging.getLogger('traffic_monitor')
|
||
|
||
GB = 1_000_000_000 # Azure 数据转移按十进制 GB 计费
|
||
|
||
|
||
def utcnow():
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def load_json(path, default):
|
||
try:
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
return json.load(f)
|
||
except FileNotFoundError:
|
||
return default
|
||
except (json.JSONDecodeError, OSError) as e:
|
||
log.warning('read %s failed: %s; using default', path, e)
|
||
return default
|
||
|
||
|
||
def save_json(path, obj):
|
||
tmp = path + '.tmp'
|
||
with open(tmp, 'w', encoding='utf-8') as f:
|
||
json.dump(obj, f, ensure_ascii=False, indent=2)
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def gb_str(n_bytes):
|
||
return f'{n_bytes / GB:.2f} GB'
|
||
|
||
|
||
def month_struct(year, month):
|
||
return {
|
||
'year': int(year),
|
||
'month': int(month),
|
||
'rx_bytes': 0,
|
||
'tx_bytes': 0,
|
||
'days': [],
|
||
'updated_at': utcnow().isoformat(),
|
||
}
|
||
|
||
|
||
class Monitor:
|
||
def __init__(self, config):
|
||
self.config = config
|
||
self.lock = threading.RLock()
|
||
self.months = {}
|
||
self.alert_state = None
|
||
|
||
os.makedirs(self.config['data_dir'], exist_ok=True)
|
||
self._load_all_months()
|
||
self._load_alert_state()
|
||
# Azure 计费用量（Usage Details API，VM 托管身份鉴权）
|
||
self.azure_usage = None
|
||
self._azure_last_fetch = 0.0
|
||
cache_path = self.config.get('azure_usage_cache_path')
|
||
if cache_path and os.path.exists(cache_path):
|
||
self.azure_usage = load_json(cache_path, None)
|
||
log.info('azure usage cache loaded: %s', self.azure_usage)
|
||
|
||
# ---------- 基础 ----------
|
||
@property
|
||
def _config_path(self):
|
||
return os.environ.get('TRAFFIC_CONFIG', '/home/agent/config.json')
|
||
|
||
def _load_all_months(self):
|
||
root = self.config['data_dir']
|
||
if not os.path.isdir(root):
|
||
return
|
||
for year_dir in os.listdir(root):
|
||
ypath = os.path.join(root, year_dir)
|
||
if not os.path.isdir(ypath):
|
||
continue
|
||
for fn in os.listdir(ypath):
|
||
if not fn.endswith('.json'):
|
||
continue
|
||
try:
|
||
with open(os.path.join(ypath, fn), 'r', encoding='utf-8') as f:
|
||
m = json.load(f)
|
||
mk = f"{m['year']:04d}-{m['month']:02d}"
|
||
self.months[mk] = m
|
||
except Exception as e:
|
||
log.warning('load %s/%s failed: %s', ypath, fn, e)
|
||
|
||
def _load_alert_state(self):
|
||
self.alert_state = load_json(
|
||
self.config['alert_state_path'],
|
||
{'alert_thresholds_sent': {}, 'monthly_reports_sent': []})
|
||
self.alert_state.setdefault('alert_thresholds_sent', {})
|
||
self.alert_state.setdefault('monthly_reports_sent', [])
|
||
|
||
def _save_alert_state(self):
|
||
save_json(self.config['alert_state_path'], self.alert_state)
|
||
|
||
def _month_path(self, year, month):
|
||
return os.path.join(self.config['data_dir'], str(year), f'{year:04d}-{month:02d}.json')
|
||
|
||
def _persist_month(self, mk):
|
||
m = self.months.get(mk)
|
||
if m is None:
|
||
return
|
||
path = self._month_path(m['year'], m['month'])
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
save_json(path, m)
|
||
|
||
# ---------- vnstat 读取 ----------
|
||
def _vnstat(self, what):
|
||
bin_path = self.config.get('vnstat_bin', 'vnstat')
|
||
iface = self.config.get('interface', 'eth0')
|
||
cmd = [bin_path, '-i', iface, '--json', what]
|
||
try:
|
||
out = subprocess.run(cmd, capture_output=True, timeout=20, check=True)
|
||
return json.loads(out.stdout.decode('utf-8'))
|
||
except FileNotFoundError:
|
||
log.error('vnstat binary not found: %s', bin_path)
|
||
return None
|
||
except subprocess.TimeoutExpired:
|
||
log.error('vnstat --json %s timed out', what)
|
||
return None
|
||
except (subprocess.CalledProcessError, json.JSONDecodeError, OSError) as e:
|
||
log.error('vnstat --json %s failed: %s', what, e)
|
||
return None
|
||
|
||
def poll_vnstat(self):
|
||
"""拉取 vnstat 月/日数据，更新并落盘。返回更新的月数。"""
|
||
monthly = self._vnstat('m')
|
||
daily = self._vnstat('d')
|
||
touched = set()
|
||
|
||
with self.lock:
|
||
if monthly:
|
||
try:
|
||
for itf in monthly.get('interfaces', []):
|
||
for item in itf.get('traffic', {}).get('month', []):
|
||
d = item.get('date', {})
|
||
mk = f"{int(d['year']):04d}-{int(d['month']):02d}"
|
||
m = self.months.get(mk)
|
||
if m is None:
|
||
m = month_struct(d['year'], d['month'])
|
||
self.months[mk] = m
|
||
m['rx_bytes'] = int(item.get('rx', 0))
|
||
m['tx_bytes'] = int(item.get('tx', 0))
|
||
m['updated_at'] = utcnow().isoformat()
|
||
touched.add(mk)
|
||
except Exception as e:
|
||
log.exception('parse vnstat monthly failed: %s', e)
|
||
|
||
if daily:
|
||
try:
|
||
by_month = {}
|
||
for itf in daily.get('interfaces', []):
|
||
for item in itf.get('traffic', {}).get('day', []):
|
||
d = item.get('date', {})
|
||
mk = f"{int(d['year']):04d}-{int(d['month']):02d}"
|
||
by_month.setdefault(mk, []).append({
|
||
'day': int(d['day']),
|
||
'rx_bytes': int(item.get('rx', 0)),
|
||
'tx_bytes': int(item.get('tx', 0)),
|
||
})
|
||
for mk, days in by_month.items():
|
||
m = self.months.get(mk)
|
||
if m is None:
|
||
y, mon = mk.split('-')
|
||
m = month_struct(int(y), int(mon))
|
||
self.months[mk] = m
|
||
days.sort(key=lambda x: x['day'])
|
||
m['days'] = days
|
||
m['updated_at'] = utcnow().isoformat()
|
||
touched.add(mk)
|
||
except Exception as e:
|
||
log.exception('parse vnstat daily failed: %s', e)
|
||
|
||
for mk in touched:
|
||
self._persist_month(mk)
|
||
return len(touched)
|
||
|
||
# ---------- 用量（按 Azure 计费周期） ----------
|
||
def current_month_key(self):
|
||
return utcnow().strftime('%Y-%m')
|
||
|
||
def cycle_start(self, now):
|
||
"""返回当前计费周期的起始 datetime（每月 billing_cycle_start_day 00:00 UTC）。"""
|
||
day = int(self.config.get('billing_cycle_start_day', 1))
|
||
day = max(1, min(28, day))
|
||
if now.day >= day:
|
||
return now.replace(day=day, hour=0, minute=0, second=0, microsecond=0)
|
||
prev = (now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||
- timedelta(days=1))
|
||
return prev.replace(day=day, hour=0, minute=0, second=0, microsecond=0)
|
||
|
||
def cycle_end(self, start):
|
||
"""周期结束 = 下个月同一天的 00:00（不含）。"""
|
||
nxt = (start.replace(day=1) + timedelta(days=32)).replace(day=1)
|
||
return nxt.replace(day=start.day, hour=0, minute=0, second=0, microsecond=0)
|
||
|
||
def _day_totals(self):
|
||
"""返回 [{date, rx, tx}]，date 为 UTC date。"""
|
||
out = []
|
||
for m in self.months.values():
|
||
for d in m.get('days', []):
|
||
out.append((date(m['year'], m['month'], d['day']),
|
||
int(d.get('rx_bytes', 0)), int(d.get('tx_bytes', 0))))
|
||
return out
|
||
|
||
def cycle_usage(self):
|
||
"""返回当前计费周期 {rx_bytes, tx_bytes, start, end}。"""
|
||
now = utcnow()
|
||
start = self.cycle_start(now)
|
||
end = self.cycle_end(start)
|
||
rx = tx = 0
|
||
with self.lock:
|
||
for dt, r, t in self._day_totals():
|
||
if start.date() <= dt <= now.date():
|
||
rx += r
|
||
tx += t
|
||
return {'rx_bytes': rx, 'tx_bytes': tx, 'start': start, 'end': end}
|
||
|
||
def tx_bytes_for(self, mk):
|
||
with self.lock:
|
||
m = self.months.get(mk, {})
|
||
return int(m.get('tx_bytes', 0))
|
||
|
||
def current_status(self):
|
||
"""当前状态：以 Azure 门户口径（Standard Data Transfer Out - Free）为主计量。"""
|
||
now = utcnow()
|
||
quota_gb = float(self.config.get('azure_free_limit_gb', 15))
|
||
cu = self.cycle_usage()
|
||
az = self.azure_usage or {}
|
||
used_gb = az.get('portal_data_transfer_out_gb')
|
||
used_pct = round(used_gb / quota_gb * 100, 4) if (used_gb is not None and quota_gb) else None
|
||
with self.lock:
|
||
cal_mk = now.strftime('%Y-%m')
|
||
m = self.months.get(cal_mk)
|
||
cal_tx = int(m.get('tx_bytes', 0)) if m else 0
|
||
cal_rx = int(m.get('rx_bytes', 0)) if m else 0
|
||
vnstat_updated = (m or {}).get('updated_at')
|
||
return {
|
||
'year': now.year,
|
||
'month': now.month,
|
||
'meter': 'azure_portal',
|
||
'used_gb': used_gb,
|
||
'quota_gb': quota_gb,
|
||
'quota_used_pct': used_pct,
|
||
'billing_cycle_start_day': int(self.config.get('billing_cycle_start_day', 1)),
|
||
'cycle_start': cu['start'].strftime('%Y-%m-%d'),
|
||
'cycle_end': cu['end'].strftime('%Y-%m-%d'),
|
||
'azure_portal_data_transfer_out_gb': az.get('portal_data_transfer_out_gb'),
|
||
'azure_bandwidth_out_gb': az.get('bandwidth_out_gb'),
|
||
'azure_usage_cost_cny': az.get('cost_cny'),
|
||
'azure_usage_month': az.get('month'),
|
||
'azure_usage_updated_at': az.get('updated_at'),
|
||
'vnstat_calendar_month_tx_bytes': cal_tx,
|
||
'vnstat_calendar_month_rx_bytes': cal_rx,
|
||
'vnstat_updated_at': vnstat_updated,
|
||
'updated_at': now.isoformat(),
|
||
}
|
||
|
||
def _vnstat_calendar_tx_gb(self):
|
||
"""参考值：vnstat 自然月出站（GB，十进制）。"""
|
||
now = utcnow()
|
||
with self.lock:
|
||
m = self.months.get(now.strftime('%Y-%m'), {})
|
||
return int(m.get('tx_bytes', 0)) / GB
|
||
|
||
# ---------- Azure 计费用量（Usage Details API） ----------
|
||
def _azure_token(self):
|
||
"""通过 VM 托管身份从 IMDS 获取 management.azure.com 的访问令牌。"""
|
||
url = ('http://169.254.169.254/metadata/identity/oauth2/token'
|
||
'?api-version=2018-02-01&resource=https%3A%2F%2Fmanagement.azure.com%2F')
|
||
req = urllib.request.Request(url, headers={'Metadata': 'true'})
|
||
with urllib.request.urlopen(req, timeout=15) as resp:
|
||
return json.loads(resp.read().decode('utf-8'))['access_token']
|
||
|
||
def _azure_subscription_id(self):
|
||
url = 'http://169.254.169.254/metadata/instance/compute?api-version=2021-02-01'
|
||
req = urllib.request.Request(url, headers={'Metadata': 'true'})
|
||
with urllib.request.urlopen(req, timeout=15) as resp:
|
||
return json.loads(resp.read().decode('utf-8')).get('subscriptionId')
|
||
|
||
def fetch_azure_usage(self, month_offset=0):
|
||
"""拉取指定月份 usageDetails，聚合与门户免费服务页同源的 Data Transfer Out 用量。
|
||
month_offset: 0=当月（写入缓存），-1=上月（用于月报，不覆盖缓存）。"""
|
||
sub = self._azure_subscription_id()
|
||
if not sub:
|
||
raise RuntimeError('cannot read subscriptionId from IMDS')
|
||
token = self._azure_token()
|
||
now = utcnow()
|
||
month_start = (now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||
+ timedelta(days=month_offset * 32)).replace(day=1)
|
||
nxt = (month_start.replace(day=1) + timedelta(days=32)).replace(day=1)
|
||
flt = ("properties/usageStart ge '{0}T00:00:00Z' and "
|
||
"properties/usageStart lt '{1}T00:00:00Z'").format(
|
||
month_start.strftime('%Y-%m-%d'), nxt.strftime('%Y-%m-%d'))
|
||
base = ('https://management.azure.com/subscriptions/{sub}/providers/'
|
||
'Microsoft.Consumption/usageDetails?api-version=2023-11-01'
|
||
'&metric=actualcost&$filter={f}')
|
||
url = base.format(sub=sub, f=urllib.parse.quote(flt, safe=''))
|
||
portal_gb = 0.0
|
||
bw_out_gb = 0.0
|
||
cost_cny = 0.0
|
||
rows = 0
|
||
pages = 0
|
||
while url:
|
||
req = urllib.request.Request(url, headers={'Authorization': 'Bearer ' + token})
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
d = json.loads(resp.read().decode('utf-8'))
|
||
for item in d.get('value', []):
|
||
p = item.get('properties', {}) or {}
|
||
qty = float(p.get('quantity') or 0)
|
||
rows += 1
|
||
name = str(p.get('meterName') or '')
|
||
subcat = str(p.get('meterSubCategory') or '')
|
||
if str(p.get('meterCategory') or '') == 'Bandwidth' and 'Data Transfer Out' in name:
|
||
bw_out_gb += qty
|
||
# 门户免费服务页 "Data Transfer Out (GB)" = MGN 路由的出站免费计费表
|
||
if 'MGN' in subcat or '- Free' in name:
|
||
portal_gb += qty
|
||
cost_cny += float(p.get('costInBillingCurrency') or 0)
|
||
pages += 1
|
||
url = d.get('nextLink')
|
||
if pages > 30:
|
||
log.warning('azure usage pagination limit reached')
|
||
break
|
||
result = {
|
||
'month': month_start.strftime('%Y-%m'),
|
||
'portal_data_transfer_out_gb': round(portal_gb, 4),
|
||
'bandwidth_out_gb': round(bw_out_gb, 4),
|
||
'cost_cny': round(cost_cny, 4),
|
||
'rows': rows,
|
||
'updated_at': utcnow().isoformat(),
|
||
}
|
||
self.azure_usage = result
|
||
if month_offset == 0:
|
||
path = self.config.get('azure_usage_cache_path')
|
||
if path:
|
||
save_json(path, result)
|
||
log.info('azure usage fetched (offset=%d): %s', month_offset, result)
|
||
return result
|
||
|
||
def check_azure_usage(self):
|
||
if not self.config.get('azure_usage_enabled', False):
|
||
return
|
||
interval = float(self.config.get('azure_usage_fetch_interval_sec', 21600))
|
||
now = time.time()
|
||
if now - self._azure_last_fetch < interval:
|
||
return
|
||
self._azure_last_fetch = now
|
||
try:
|
||
self.fetch_azure_usage()
|
||
except Exception as e:
|
||
log.warning('azure usage fetch failed: %s', e)
|
||
|
||
# ---------- IP / 邮件 ----------
|
||
def get_current_ip(self):
|
||
configured = self.config.get('current_ip')
|
||
if configured:
|
||
return str(configured)
|
||
try:
|
||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
try:
|
||
s.connect(('8.8.8.8', 80))
|
||
return s.getsockname()[0]
|
||
finally:
|
||
s.close()
|
||
except OSError:
|
||
return '0.0.0.0'
|
||
|
||
def send_email(self, subject, body):
|
||
cfg = self.config['smtp']
|
||
msg = EmailMessage()
|
||
msg['Subject'] = subject
|
||
msg['From'] = formataddr((cfg.get('from_name', ''), cfg.get('from_addr')))
|
||
msg['To'] = cfg.get('to_addr')
|
||
msg.set_content(body)
|
||
|
||
context = None
|
||
if cfg.get('skip_cert_verify'):
|
||
context = ssl._create_unverified_context()
|
||
host = cfg['host']
|
||
port = int(cfg.get('port', 465))
|
||
timeout = int(cfg.get('timeout_sec', 30))
|
||
if cfg.get('ssl', False):
|
||
smtp_cls = smtplib.SMTP_SSL
|
||
kwargs = {'context': context}
|
||
else:
|
||
smtp_cls = smtplib.SMTP
|
||
kwargs = {}
|
||
with smtp_cls(host, port, timeout=timeout, **kwargs) as smtp:
|
||
if not cfg.get('ssl', False):
|
||
smtp.ehlo()
|
||
if cfg.get('starttls', True):
|
||
smtp.starttls(context=context)
|
||
smtp.ehlo()
|
||
smtp.login(cfg['user'], cfg['password'])
|
||
smtp.send_message(msg)
|
||
log.info('email sent: %s', subject)
|
||
|
||
def build_alert_mail(self, mk, used_gb, quota_gb, pct, az):
|
||
ip = self.get_current_ip()
|
||
month = (az or {}).get('month', mk.replace('azure:', ''))
|
||
az_updated = (az or {}).get('updated_at', '')
|
||
vnstat_tx = self._vnstat_calendar_tx_gb()
|
||
body = (
|
||
'Royirin Rosmontis Private Proxy JP 流量监控\n'
|
||
'==========\n'
|
||
'计量口径：Azure 门户（Standard Data Transfer Out - Free）\n'
|
||
f'统计月份：{month}（自然月）\n'
|
||
f'当前用量：{used_gb:.2f} GB / {quota_gb:.0f} GB（{pct:.2f}%）\n'
|
||
f'AZURE 数据更新：{az_updated}\n'
|
||
f'参考(vnstat 自然月出站)：{vnstat_tx:.2f} GB\n'
|
||
f'当前 IP：{ip}\n'
|
||
f'生效域名：{self.config.get("effective_domain", "")}\n'
|
||
'==========\n'
|
||
f'已达阈值 {pct:.2f}% (Azure 门户口径)，请注意使用情况'
|
||
)
|
||
subject = f'[流量告警] {month} 已用 {pct:.2f}% ({used_gb:.2f}/{quota_gb:.0f} GB)'
|
||
return subject, body
|
||
|
||
def build_monthly_mail(self, month, used_gb, quota_gb, pct, az):
|
||
ip = self.get_current_ip()
|
||
over = '是' if used_gb > quota_gb else '否'
|
||
az_updated = (az or {}).get('updated_at', '')
|
||
vnstat_tx = self._vnstat_calendar_tx_gb()
|
||
body = (
|
||
'Royirin Rosmontis Private Proxy JP 流量监控\n'
|
||
'月度报告\n'
|
||
'==========\n'
|
||
'计量口径：Azure 门户（Standard Data Transfer Out - Free）\n'
|
||
f'报告月份：{month}（自然月）\n'
|
||
f'当月用量：{used_gb:.2f} GB / {quota_gb:.0f} GB（{pct:.2f}%）\n'
|
||
f'AZURE 数据更新：{az_updated}\n'
|
||
f'参考(vnstat 自然月出站)：{vnstat_tx:.2f} GB\n'
|
||
f'当前 IP：{ip}\n'
|
||
f'生效域名：{self.config.get("effective_domain", "")}\n'
|
||
'==========\n'
|
||
f'上月使用 {over} 超限, 使用量为限额的 {pct:.2f}%'
|
||
)
|
||
subject = f'[流量月报] {month} 用量 {used_gb:.2f} GB ({pct:.2f}%)'
|
||
return subject, body
|
||
|
||
# ---------- 定时任务 ----------
|
||
def check_alerts(self):
|
||
"""阈值告警：以 Azure 门户口径为准（无 Azure 数据时跳过，避免误报）。"""
|
||
quota_gb = float(self.config.get('azure_free_limit_gb', 15))
|
||
if quota_gb <= 0:
|
||
return
|
||
az = self.azure_usage or {}
|
||
used_gb = az.get('portal_data_transfer_out_gb')
|
||
if used_gb is None:
|
||
log.warning('alerts skipped: no azure usage data yet')
|
||
return
|
||
mk = 'azure:' + str(az.get('month') or utcnow().strftime('%Y-%m'))
|
||
pct = used_gb / quota_gb * 100
|
||
sent = self.alert_state['alert_thresholds_sent'].setdefault(mk, [])
|
||
if not sent:
|
||
# 新周期/口径切换：已越过的档位直接标记，避免切换瞬间补发一堆邮件
|
||
for t in sorted(self.config.get('alert_thresholds_pct', [])):
|
||
if pct >= t:
|
||
sent.append(t)
|
||
self._save_alert_state()
|
||
log.info('seeded thresholds for %s at %.2f%%: %s', mk, pct, sent)
|
||
for t in sorted(self.config.get('alert_thresholds_pct', [])):
|
||
if pct >= t and t not in sent:
|
||
subject, body = self.build_alert_mail(mk, used_gb, quota_gb, pct, az)
|
||
try:
|
||
self.send_email(subject, body)
|
||
sent.append(t)
|
||
self._save_alert_state()
|
||
log.info('threshold %s%% alert sent for %s (used %.2f%%)', t, mk, pct)
|
||
except Exception as e:
|
||
log.error('send threshold %s%% alert failed for %s: %s; will retry',
|
||
t, mk, e)
|
||
break
|
||
|
||
def check_monthly_report(self):
|
||
"""月度报告：以 Azure 门户口径为准（上月数据，当月 1 号后发送）。"""
|
||
now = utcnow()
|
||
cur_start = self.cycle_start(now)
|
||
prev_start = (cur_start.replace(day=1) - timedelta(days=1)).replace(day=1)
|
||
prev_mk = 'azure:' + prev_start.strftime('%Y-%m')
|
||
if prev_mk in self.alert_state['monthly_reports_sent']:
|
||
return
|
||
# 口径切换迁移：旧键(cycle:)已发过则直接标记，避免重复发信
|
||
old_mk = 'cycle:' + prev_start.strftime('%Y-%m')
|
||
if old_mk in self.alert_state['monthly_reports_sent']:
|
||
self.alert_state['monthly_reports_sent'].append(prev_mk)
|
||
self._save_alert_state()
|
||
log.info('monthly report key migrated: %s already sent as %s', prev_mk, old_mk)
|
||
return
|
||
try:
|
||
az_prev = self.fetch_azure_usage(month_offset=-1)
|
||
except Exception as e:
|
||
log.error('fetch prev month azure usage failed: %s; will retry', e)
|
||
return
|
||
quota_gb = float(self.config.get('azure_free_limit_gb', 15))
|
||
used_gb = az_prev.get('portal_data_transfer_out_gb') or 0.0
|
||
pct = used_gb / quota_gb * 100 if quota_gb else 0.0
|
||
subject, body = self.build_monthly_mail(
|
||
az_prev.get('month', prev_start.strftime('%Y-%m')), used_gb, quota_gb, pct, az_prev)
|
||
try:
|
||
self.send_email(subject, body)
|
||
self.alert_state['monthly_reports_sent'].append(prev_mk)
|
||
self._save_alert_state()
|
||
log.info('monthly report sent for %s (used %.2f GB, %.2f%%)',
|
||
prev_mk, used_gb, pct)
|
||
except Exception as e:
|
||
log.error('send monthly report for %s failed: %s; will retry', prev_mk, e)
|
||
|
||
|
||
def build_http_handler(monitor):
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def do_GET(self):
|
||
if self.path == '/vpn/status/traffic':
|
||
try:
|
||
payload = json.dumps(monitor.current_status(), ensure_ascii=False)
|
||
body = payload.encode('utf-8')
|
||
self.send_response(200)
|
||
self.send_header('Content-Type', 'application/json; charset=utf-8')
|
||
self.send_header('Content-Length', str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
except Exception as e:
|
||
log.exception('status endpoint error')
|
||
body = b'{"error":"internal"}'
|
||
self.send_response(500)
|
||
self.send_header('Content-Type', 'application/json')
|
||
self.send_header('Content-Length', str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
elif self.path in ('/', '/healthz'):
|
||
body = b'ok\n'
|
||
self.send_response(200)
|
||
self.send_header('Content-Type', 'text/plain')
|
||
self.send_header('Content-Length', str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
else:
|
||
body = b'{"error":"not found"}\n'
|
||
self.send_response(404)
|
||
self.send_header('Content-Type', 'application/json')
|
||
self.send_header('Content-Length', str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def log_message(self, fmt, *args):
|
||
log.debug('http: ' + fmt, *args)
|
||
|
||
return Handler
|
||
|
||
|
||
def run_service(config):
|
||
monitor = Monitor(config)
|
||
|
||
def poll_loop():
|
||
while True:
|
||
try:
|
||
n = monitor.poll_vnstat()
|
||
log.debug('vnstat poll updated %d month(s)', n)
|
||
except Exception:
|
||
log.exception('poll loop error')
|
||
time.sleep(int(config.get('poll_interval_sec', 60)))
|
||
|
||
def schedule_loop():
|
||
while True:
|
||
try:
|
||
monitor.check_azure_usage() # 先刷新 Azure 口径数据，再评估告警
|
||
monitor.check_alerts()
|
||
monitor.check_monthly_report()
|
||
except Exception:
|
||
log.exception('schedule loop error')
|
||
time.sleep(int(config.get('schedule_interval_sec', 30)))
|
||
|
||
threading.Thread(target=poll_loop, daemon=True, name='vnstat-poll').start()
|
||
threading.Thread(target=schedule_loop, daemon=True, name='scheduler').start()
|
||
|
||
host = config.get('http_host', '127.0.0.1')
|
||
port = int(config.get('http_port', 18080))
|
||
server = ThreadingHTTPServer((host, port), build_http_handler(monitor))
|
||
log.info('HTTP server listening on http://%s:%d/vpn/status/traffic', host, port)
|
||
server.serve_forever()
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description='vnstat traffic monitor')
|
||
ap.add_argument('--config', default=os.environ.get('TRAFFIC_CONFIG', '/home/agent/config.json'))
|
||
ap.add_argument('--once', action='store_true', help='拉取一次 vnstat 后退出')
|
||
ap.add_argument('--status', action='store_true', help='打印当前状态 JSON 后退出')
|
||
ap.add_argument('--test-email', action='store_true', help='发送一封测试邮件后退出')
|
||
args = ap.parse_args()
|
||
|
||
os.environ['TRAFFIC_CONFIG'] = args.config
|
||
with open(args.config, 'r', encoding='utf-8') as f:
|
||
config = json.load(f)
|
||
|
||
logging.basicConfig(
|
||
level=getattr(logging, config.get('log_level', 'INFO').upper(), logging.INFO),
|
||
format='%(asctime)s %(levelname)s %(name)s: %(message)s')
|
||
|
||
if args.once:
|
||
m = Monitor(config)
|
||
n = m.poll_vnstat()
|
||
print(f'updated {n} month(s)')
|
||
return 0
|
||
if args.status:
|
||
m = Monitor(config)
|
||
m.poll_vnstat()
|
||
print(json.dumps(m.current_status(), ensure_ascii=False, indent=2))
|
||
return 0
|
||
if args.test_email:
|
||
m = Monitor(config)
|
||
m.poll_vnstat()
|
||
quota_gb = float(config.get('azure_free_limit_gb', 15))
|
||
az = m.azure_usage or {}
|
||
used_gb = az.get('portal_data_transfer_out_gb')
|
||
if used_gb is None:
|
||
try:
|
||
m.fetch_azure_usage()
|
||
az = m.azure_usage or {}
|
||
used_gb = az.get('portal_data_transfer_out_gb')
|
||
except Exception as e:
|
||
print('azure usage fetch failed:', e)
|
||
mk = 'azure:' + str(az.get('month') or utcnow().strftime('%Y-%m'))
|
||
pct = used_gb / quota_gb * 100 if (used_gb is not None and quota_gb) else 0.0
|
||
subject, body = m.build_alert_mail(mk, used_gb or 0.0, quota_gb, pct, az)
|
||
print('sending test email:')
|
||
print('Subject:', subject)
|
||
print(body)
|
||
m.send_email(subject, body)
|
||
print('sent')
|
||
return 0
|
||
run_service(config)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
sys.exit(main())
|