wgctl/core/lib/peers.py

"""
peers.py — peer data, transfer stats, group lookups.
"""

import os
import json
import sys
import glob

from lib.util import DATETIME_FMT, build_ip_to_name, build_pubkey_to_name, fmt_ts


def peer_data(clients_dir, meta_dir, events_log):
    """
    Output: name|ip|rule|type|last_ts|last_evt|main_group
    """
    meta = {}
    for f in glob.glob(f"{meta_dir}/*.meta"):
        name = os.path.basename(f).replace('.meta', '')
        try:
            with open(f) as mf:
                meta[name] = json.load(mf)
        except Exception:
            meta[name] = {}

    last_events = {}
    try:
        with open(events_log) as f:
            for line in f:
                try:
                    e = json.loads(line.strip())
                    client = e.get('client', '')
                    if client:
                        last_events[client] = e
                except Exception:
                    pass
    except Exception:
        pass

    for conf in sorted(glob.glob(f"{clients_dir}/*.conf")):
        name = os.path.basename(conf).replace('.conf', '')
        ip = ''
        try:
            with open(conf) as f:
                for line in f:
                    if line.startswith('Address'):
                        ip = line.split('=')[1].strip().split('/')[0]
                        break
        except Exception:
            pass

        m          = meta.get(name, {})
        rule       = m.get('rule', '')
        peer_type  = m.get('type', '')
        main_group = m.get('main_group', '')

        last_event = last_events.get(name, {})
        last_ts    = last_event.get('timestamp', '')
        last_evt   = last_event.get('event', '')

        print(f"{name}|{ip}|{rule}|{peer_type}|{last_ts}|{last_evt}|{main_group}")


def peer_transfer(wg_interface):
    """Get total transfer bytes per peer."""
    import subprocess
    low  = int(os.environ.get('ACTIVITY_TOTAL_LOW_BYTES',  '1000000'))
    med  = int(os.environ.get('ACTIVITY_TOTAL_MED_BYTES',  '10000000'))
    high = int(os.environ.get('ACTIVITY_TOTAL_HIGH_BYTES', '100000000'))
    try:
        result = subprocess.run(
            ['wg', 'show', wg_interface, 'transfer'],
            capture_output=True, text=True
        )
        for line in result.stdout.strip().split('\n'):
            if not line:
                continue
            parts = line.split('\t')
            if len(parts) == 3:
                pubkey, rx, tx = parts
                total = int(rx) + int(tx)
                if total == 0:        level = 'none'
                elif total < low:     level = 'low'
                elif total < med:     level = 'medium'
                elif total < high:    level = 'high'
                else:                 level = 'very high'
                print(f"{pubkey}|{rx}|{tx}|{level}")
    except Exception:
        pass


def peer_transfer_delta(wg_interface, cache_file):
    """Calculate current transfer rate using delta from previous sample."""
    import subprocess
    import time

    low  = int(os.environ.get('ACTIVITY_CURRENT_LOW_BYTES',  '10000'))
    med  = int(os.environ.get('ACTIVITY_CURRENT_MED_BYTES',  '100000'))
    high = int(os.environ.get('ACTIVITY_CURRENT_HIGH_BYTES', '1000000'))

    current = {}
    now = time.time()
    try:
        result = subprocess.run(
            ['wg', 'show', wg_interface, 'transfer'],
            capture_output=True, text=True
        )
        for line in result.stdout.strip().split('\n'):
            if not line:
                continue
            parts = line.split('\t')
            if len(parts) == 3:
                pubkey, rx, tx = parts
                current[pubkey] = {'rx': int(rx), 'tx': int(tx), 'ts': now}
    except Exception:
        pass

    prev = {}
    if os.path.exists(cache_file):
        try:
            with open(cache_file) as f:
                prev = json.load(f)
        except Exception:
            pass

    try:
        with open(cache_file, 'w') as f:
            json.dump(current, f)
    except Exception:
        pass

    for pubkey, data in current.items():
        if pubkey in prev:
            dt = data['ts'] - prev[pubkey].get('ts', data['ts'])
            if dt > 0:
                rx_rate = max(0, (data['rx'] - prev[pubkey]['rx']) / dt)
                tx_rate = max(0, (data['tx'] - prev[pubkey]['tx']) / dt)
                total   = rx_rate + tx_rate
                if total <= 0:        level = 'idle'
                elif total < low:     level = 'low'
                elif total < med:     level = 'medium'
                elif total < high:    level = 'high'
                else:                 level = 'very high'
                print(f"{pubkey}|{int(rx_rate)}|{int(tx_rate)}|{level}")
            else:
                print(f"{pubkey}|0|0|idle")
        else:
            print(f"{pubkey}|0|0|unknown")


def peer_group_map(groups_dir):
    """Return peer:group pairs for all groups."""
    try:
        for group_file in glob.glob(f"{groups_dir}/*.group"):
            try:
                with open(group_file) as f:
                    g = json.load(f)
                name = g.get('name', '')
                for peer in g.get('peers', []):
                    if peer:
                        print(f"{peer}:{name}")
            except Exception:
                pass
    except Exception:
        pass


def peer_groups(groups_dir, peer_name):
    """Find all groups containing a peer."""
    try:
        for group_file in glob.glob(f"{groups_dir}/*.group"):
            try:
                with open(group_file) as f:
                    g = json.load(f)
                if peer_name in g.get('peers', []):
                    print(g.get('name', ''))
            except Exception:
                pass
    except Exception:
        pass