pic/api/wireguard_manager.py

#!/usr/bin/env python3
"""
WireGuard Manager for Personal Internet Cell
"""

import os
import json
import base64
import subprocess
import logging
from datetime import datetime
from typing import Dict, List, Optional, Any
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from base_service_manager import BaseServiceManager

logger = logging.getLogger(__name__)

SERVER_ADDRESS = '172.20.0.1/16'
SERVER_NETWORK = '172.20.0.0/16'
PEER_DNS = '172.20.0.2'
DEFAULT_PORT = 51820


class WireGuardManager(BaseServiceManager):
    """Manages WireGuard VPN configuration and peers"""

    def __init__(self, data_dir: str = '/app/data', config_dir: str = '/app/config'):
        super().__init__('wireguard', data_dir, config_dir)
        self.wireguard_dir = os.path.join(config_dir, 'wireguard')
        self.keys_dir = os.path.join(data_dir, 'wireguard', 'keys')
        self.peers_dir = os.path.join(data_dir, 'wireguard', 'peers')

        self.safe_makedirs(self.wireguard_dir)
        self.safe_makedirs(self.keys_dir)
        self.safe_makedirs(os.path.join(self.keys_dir, 'peers'))
        self.safe_makedirs(self.peers_dir)

        self._ensure_server_keys()

    # ── Key management ────────────────────────────────────────────────────────

    @staticmethod
    def _generate_keypair():
        """Return (private_bytes, public_bytes) using X25519."""
        priv = X25519PrivateKey.generate()
        return priv.private_bytes_raw(), priv.public_key().public_bytes_raw()

    def _ensure_server_keys(self):
        priv_file = os.path.join(self.keys_dir, 'private.key')
        pub_file = os.path.join(self.keys_dir, 'public.key')
        if not os.path.exists(priv_file):
            try:
                priv_bytes, pub_bytes = self._generate_keypair()
                with open(priv_file, 'wb') as f:
                    f.write(priv_bytes)
                with open(pub_file, 'wb') as f:
                    f.write(pub_bytes)
            except (PermissionError, OSError):
                pass

    def get_keys(self) -> Dict[str, str]:
        """Return server public/private keys as base64 strings."""
        priv_file = os.path.join(self.keys_dir, 'private.key')
        pub_file = os.path.join(self.keys_dir, 'public.key')
        with open(priv_file, 'rb') as f:
            priv = f.read()
        with open(pub_file, 'rb') as f:
            pub = f.read()
        return {
            'private_key': base64.b64encode(priv).decode(),
            'public_key': base64.b64encode(pub).decode(),
        }

    def generate_peer_keys(self, peer_name: str) -> Dict[str, str]:
        """Generate a keypair for a peer, save to keys_dir/peers/, return as base64."""
        priv_bytes, pub_bytes = self._generate_keypair()
        priv_b64 = base64.b64encode(priv_bytes).decode()
        pub_b64 = base64.b64encode(pub_bytes).decode()

        peer_keys_dir = os.path.join(self.keys_dir, 'peers')
        with open(os.path.join(peer_keys_dir, f'{peer_name}_private.key'), 'w') as f:
            f.write(priv_b64)
        with open(os.path.join(peer_keys_dir, f'{peer_name}_public.key'), 'w') as f:
            f.write(pub_b64)

        return {'private_key': priv_b64, 'public_key': pub_b64, 'peer_name': peer_name}

    # ── Config generation ─────────────────────────────────────────────────────

    def get_config(self, interface: str = 'wg0', port: int = DEFAULT_PORT):
        """Return server config (alias for generate_config, returns dict for API compat)."""
        return {'config': self.generate_config(interface, port)}

    def generate_config(self, interface: str = 'wg0', port: int = DEFAULT_PORT) -> str:
        """Return a WireGuard [Interface] config string for the server."""
        keys = self.get_keys()
        return (
            f'[Interface]\n'
            f'PrivateKey = {keys["private_key"]}\n'
            f'Address = {SERVER_ADDRESS}\n'
            f'ListenPort = {port}\n'
            f'PostUp = iptables -A FORWARD -i %i -j ACCEPT; '
            f'iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE\n'
            f'PostDown = iptables -D FORWARD -i %i -j ACCEPT; '
            f'iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE\n'
        )

    def _config_file(self) -> str:
        return os.path.join(self.wireguard_dir, 'wg0.conf')

    def _read_config(self) -> str:
        cf = self._config_file()
        if os.path.exists(cf):
            with open(cf, 'r') as f:
                return f.read()
        return self.generate_config()

    def _write_config(self, content: str):
        with open(self._config_file(), 'w') as f:
            f.write(content)

    # ── Peer CRUD ─────────────────────────────────────────────────────────────

    def add_peer(self, name: str, public_key: str, endpoint_ip: str,
                 allowed_ips: str = SERVER_NETWORK,
                 persistent_keepalive: int = 25) -> bool:
        """Add a [Peer] block to wg0.conf."""
        try:
            content = self._read_config()
            peer_block = (
                f'\n[Peer]\n'
                f'# {name}\n'
                f'PublicKey = {public_key}\n'
                f'AllowedIPs = {allowed_ips}\n'
                f'PersistentKeepalive = {persistent_keepalive}\n'
            )
            if endpoint_ip:
                peer_block += f'Endpoint = {endpoint_ip}:{DEFAULT_PORT}\n'
            self._write_config(content + peer_block)
            return True
        except Exception as e:
            logger.error(f'add_peer failed: {e}')
            return False

    def remove_peer(self, public_key: str) -> bool:
        """Remove the [Peer] block matching public_key from wg0.conf."""
        try:
            content = self._read_config()
            # Split on blank lines between blocks
            raw_blocks = ('\n' + content).split('\n\n')
            new_blocks = [
                b for b in raw_blocks
                if not (f'PublicKey = {public_key}' in b and '[Peer]' in b)
            ]
            self._write_config('\n\n'.join(new_blocks).lstrip('\n'))
            return True
        except Exception as e:
            logger.error(f'remove_peer failed: {e}')
            return False

    def get_peers(self) -> List[Dict[str, Any]]:
        """Parse wg0.conf and return list of peer dicts."""
        content = self._read_config()
        peers = []
        sections = content.split('[Peer]')
        for section in sections[1:]:
            peer: Dict[str, Any] = {}
            for line in section.strip().splitlines():
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                if '=' not in line:
                    continue
                key, _, value = line.partition('=')
                key = key.strip().lower().replace(' ', '')
                value = value.strip()
                if key == 'publickey':
                    peer['public_key'] = value
                elif key == 'allowedips':
                    peer['allowed_ips'] = value
                elif key == 'persistentkeepalive':
                    try:
                        peer['persistent_keepalive'] = int(value)
                    except ValueError:
                        peer['persistent_keepalive'] = value
                elif key == 'endpoint':
                    peer['endpoint'] = value
            if peer:
                peers.append(peer)
        return peers

    def update_peer_ip(self, public_key: str, new_ip: str) -> bool:
        """Update AllowedIPs for the peer with the given public key."""
        content = self._read_config()
        if f'PublicKey = {public_key}' not in content:
            return False
        lines = content.splitlines()
        in_target = False
        new_lines = []
        for line in lines:
            if line.strip() == f'PublicKey = {public_key}':
                in_target = True
            if in_target and line.strip().startswith('AllowedIPs'):
                line = f'AllowedIPs = {new_ip}'
                in_target = False
            new_lines.append(line)
        self._write_config('\n'.join(new_lines))
        return True

    def get_peer_config(self, peer_name: str, peer_ip: str,
                        peer_private_key: str,
                        server_endpoint: str = '<SERVER_IP>') -> str:
        """Generate a WireGuard client config string."""
        server_keys = self.get_keys()
        return (
            f'[Interface]\n'
            f'PrivateKey = {peer_private_key}\n'
            f'Address = {peer_ip}/32\n'
            f'DNS = {PEER_DNS}\n'
            f'\n'
            f'[Peer]\n'
            f'PublicKey = {server_keys["public_key"]}\n'
            f'AllowedIPs = {SERVER_NETWORK}\n'
            f'Endpoint = {server_endpoint}:{DEFAULT_PORT}\n'
            f'PersistentKeepalive = 25\n'
        )

    # ── Status & connectivity ─────────────────────────────────────────────────

    def get_status(self) -> Dict[str, Any]:
        """Return service status by checking whether the Docker container is up."""
        try:
            result = subprocess.run(
                ['docker', 'ps', '--filter', 'name=cell-wireguard', '--format', '{{.Names}}'],
                capture_output=True, text=True, timeout=5,
            )
            running = 'cell-wireguard' in result.stdout
            return {
                'running': running,
                'status': 'online' if running else 'offline',
                'interface': 'wg0',
                'ip_info': {'address': SERVER_ADDRESS} if running else {},
                'peers_count': len(self.get_peers()),
                'timestamp': datetime.utcnow().isoformat(),
            }
        except Exception as e:
            return self.handle_error(e, 'get_status')

    def test_connectivity(self, peer_ip: str) -> Dict[str, Any]:
        """Ping a peer IP and return results."""
        try:
            result = subprocess.run(
                ['ping', '-c', '1', '-W', '2', peer_ip],
                capture_output=True, text=True, timeout=5,
            )
            return {
                'peer_ip': peer_ip,
                'ping_success': result.returncode == 0,
                'ping_output': result.stdout,
                'ping_error': result.stderr,
            }
        except Exception as e:
            return {
                'peer_ip': peer_ip,
                'ping_success': False,
                'ping_output': '',
                'ping_error': str(e),
            }

    def get_metrics(self) -> Dict[str, Any]:
        status = self.get_status()
        return {
            'service': 'wireguard',
            'timestamp': datetime.utcnow().isoformat(),
            'status': status.get('status', 'unknown'),
            'peers_count': status.get('peers_count', 0),
        }

    def restart_service(self) -> bool:
        try:
            result = subprocess.run(
                ['docker', 'restart', 'cell-wireguard'],
                capture_output=True, text=True, timeout=30,
            )
            return result.returncode == 0
        except Exception as e:
            logger.error(f'restart_service failed: {e}')
            return False