Сборка статистики по использованию сети

execute_statistics.py

#!/root/audit_env/bin/python

import psycopg2
import psycopg2.extras

import sqlite3
import socket
import json
import pathlib
import urllib.request
from typing import Dict

CONFIG_PATH = "/root/server_config.json"
POSTGRES_CONFIG_PATH = "/root/pg_config.json"
VNSTAT_DB_PATH = "/var/lib/vnstat/vnstat.db"

SQL_QUERY = """
select 
    date
    , rx
    , tx
    -- , round((rx + tx) / 1024.0 / 1024 / 1024, 2) as total -- Total Bytes 
    -- , round((rx + tx) * 8.0 / 1000000 / (5 * 60), 2) as avg_speed -- avg speed Mbit/s
from fiveminute 
where interface = 1
order by date desc
"""

INSERT_TEMPLATE = """
    insert into dwh.dt_server_data({columns}) values({placeholders})
    on conflict (ip) do update set
    rx = excluded.rx,
    tx = excluded.tx,
    vnstat_date = excluded.vnstat_date
""".strip()


def get_postgres_config() -> Dict:
    if pathlib.Path(POSTGRES_CONFIG_PATH).exists():
        with open(POSTGRES_CONFIG_PATH, "r") as f:
            return json.loads(f.read())
    else:
        return {}


def get_or_create_base_config() -> Dict:
    if pathlib.Path(CONFIG_PATH).exists():
        with open(CONFIG_PATH, "r") as f:
            return json.loads(f.read())

    public_ip = urllib.request.urlopen('https://api.ipify.org').read().decode('utf8')
    public_ip_info = urllib.request.urlopen(f'http://ip-api.com/json/{public_ip}')

    data = json.load(public_ip_info)
    country = data.get('country', 'Unknown')
    provider = data.get('org', 'Unknown')
    city = data.get('city', 'Unknown')
    server_name = socket.gethostname()

    raw_data = {
        "ip": public_ip,
        "country": country,
        "provider": provider,
        "city": city,
        "server_name": server_name,
    }

    with open(CONFIG_PATH, "w") as f:
        f.write(json.dumps(raw_data))
    print(f"File: {CONFIG_PATH} not found. Create with data: {raw_data}")
    return raw_data


def fetch_network_dump(cursor) -> Dict:
    columns = ["vnstat_date", "rx", "tx"]
    cursor.execute(SQL_QUERY.strip())
    data = cursor.fetchone()
    return dict(zip(columns, data))


def collect_server_data() -> Dict:
    final_data = {}

    conn = sqlite3.connect(VNSTAT_DB_PATH)
    cursor = conn.cursor()
    config_data = get_or_create_base_config()
    vnstat_data = fetch_network_dump(cursor)
    final_data.update(config_data)
    final_data.update(vnstat_data)

    return final_data


def insert_stats(data: Dict) -> None:
    columns = ", ".join(data.keys())
    placeholders = ", ".join(["%s" for _ in range(len(data.keys()))])

    cmd = INSERT_TEMPLATE.format(
        columns=columns, placeholders=placeholders
    )
    pg_cursor.execute(cmd, tuple(data.values()))
    print(f"Inserted success: {data}")


pg_config = get_postgres_config()
pg_conn = psycopg2.connect(
    dbname=pg_config["dbname"],
    user=pg_config["user"],
    password=pg_config["password"],
    host=pg_config["host"],
    port=5432
)
pg_cursor = pg_conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
pg_conn.autocommit = True

server_statistics = collect_server_data()
insert_stats(server_statistics)

pg_conn.close()