pic/api/routes/peers.py

import logging
import ipaddress
from flask import Blueprint, request, jsonify, session
logger = logging.getLogger('picell')
bp = Blueprint('peers', __name__)


def _next_peer_ip() -> str:
    """Auto-assign the next free host address from the configured VPN subnet."""
    from app import wireguard_manager, peer_registry
    server_addr = wireguard_manager._get_configured_address()
    network = ipaddress.ip_network(server_addr, strict=False)
    server_ip = str(ipaddress.ip_interface(server_addr).ip)
    used = {p.get('ip', '').split('/')[0] for p in peer_registry.list_peers()}
    for host in network.hosts():
        ip = str(host)
        if ip == server_ip:
            continue
        if ip not in used:
            return ip
    raise ValueError(f'No free IPs left in {network}')


@bp.route('/api/peers', methods=['GET'])
def get_peers():
    try:
        from app import peer_registry
        return jsonify(peer_registry.list_peers())
    except Exception as e:
        logger.error(f"Error getting peers: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers', methods=['POST'])
def add_peer():
    """Add a peer and auto-provision auth/email/calendar/files accounts."""
    try:
        from app import (peer_registry, wireguard_manager, firewall_manager,
                         email_manager, calendar_manager, file_manager, auth_manager,
                         cell_link_manager, _configured_domain, COREFILE_PATH)
        try:
            _wg_addr = wireguard_manager._get_configured_address()
            _wg_subnet = str(ipaddress.ip_network(_wg_addr, strict=False)) if _wg_addr else '10.0.0.0/24'
        except Exception:
            _wg_subnet = '10.0.0.0/24'
        _cell_links = cell_link_manager.list_connections()
        _cell_subnets = [l['vpn_subnet'] for l in _cell_links if l.get('vpn_subnet')]
        data = request.get_json(silent=True)
        if data is None:
            return jsonify({"error": "No data provided"}), 400

        for field in ('name', 'public_key'):
            if field not in data:
                return jsonify({"error": f"Missing required field: {field}"}), 400

        password = data.get('password') or ''
        if not password:
            return jsonify({"error": "Missing required field: password"}), 400
        if len(password) < 10:
            return jsonify({"error": "password must be at least 10 characters"}), 400

        try:
            assigned_ip = data.get('ip') or _next_peer_ip()
        except ValueError as e:
            return jsonify({'error': str(e)}), 409

        _valid_services = {'calendar', 'files', 'mail', 'webdav'}
        service_access = data.get('service_access', list(_valid_services))
        if not isinstance(service_access, list) or not all(s in _valid_services for s in service_access):
            return jsonify({"error": f"service_access must be a list of: {sorted(_valid_services)}"}), 400

        peer_name = data['name']

        if not auth_manager.create_user(peer_name, password, 'peer'):
            return jsonify({"error": "Could not create auth account (duplicate name?)"}), 400

        provisioned = ['auth']
        domain = _configured_domain()
        for step_name, step_fn in [
            ('email',    lambda: email_manager.create_email_user(peer_name, domain, password)),
            ('calendar', lambda: calendar_manager.create_calendar_user(peer_name, password)),
            ('files',    lambda: file_manager.create_user(peer_name, password)),
        ]:
            try:
                if step_fn():
                    provisioned.append(step_name)
                else:
                    logger.warning(f"Peer {peer_name}: {step_name} account creation returned False")
            except Exception as e:
                logger.warning(f"Peer {peer_name}: {step_name} account creation failed (non-fatal): {e}")

        peer_info = {
            'peer': peer_name,
            'ip': assigned_ip,
            'public_key': data['public_key'],
            'private_key': data.get('private_key'),
            'server_public_key': data.get('server_public_key'),
            'server_endpoint': data.get('server_endpoint'),
            'allowed_ips': data.get('allowed_ips'),
            'persistent_keepalive': data.get('persistent_keepalive'),
            'description': data.get('description'),
            'internet_access': data.get('internet_access', True),
            'service_access': service_access,
            'peer_access': data.get('peer_access', True),
            'config_needs_reinstall': False,
        }

        peer_added_to_registry = False
        firewall_applied = False
        try:
            success = peer_registry.add_peer(peer_info)
            if not success:
                for svc in ('files', 'calendar', 'email', 'auth'):
                    try:
                        if svc == 'files':
                            file_manager.delete_user(peer_name)
                        elif svc == 'calendar':
                            calendar_manager.delete_calendar_user(peer_name)
                        elif svc == 'email':
                            email_manager.delete_email_user(peer_name, _configured_domain())
                        elif svc == 'auth':
                            auth_manager.delete_user(peer_name)
                    except Exception:
                        pass
                return jsonify({"error": f"Peer {peer_name} already exists"}), 400
            peer_added_to_registry = True

            firewall_manager.apply_peer_rules(peer_info['ip'], peer_info,
                                              wg_subnet=_wg_subnet, cell_subnets=_cell_subnets)
            firewall_applied = True

            wg_allowed = f"{assigned_ip}/32" if '/' not in assigned_ip else assigned_ip
            try:
                wireguard_manager.add_peer(peer_name, data['public_key'], endpoint_ip='', allowed_ips=wg_allowed)
            except Exception as wg_err:
                logger.warning(f"Peer {peer_name}: WireGuard server config update failed (non-fatal): {wg_err}")

            firewall_manager.apply_all_dns_rules(peer_registry.list_peers(), COREFILE_PATH, _configured_domain(),
                                                 cell_links=cell_link_manager.list_connections())
            return jsonify({"message": f"Peer {peer_name} added successfully", "ip": assigned_ip}), 201

        except Exception as e:
            if firewall_applied:
                try:
                    firewall_manager.clear_peer_rules(peer_info['ip'])
                except Exception:
                    pass
            if peer_added_to_registry:
                try:
                    peer_registry.remove_peer(peer_name)
                except Exception:
                    pass
            logger.error(f"Error adding peer {peer_name}: {e}")
            return jsonify({'error': str(e)}), 500

    except Exception as e:
        logger.error(f"Error adding peer: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/<peer_name>', methods=['PUT'])
def update_peer(peer_name):
    try:
        from app import (peer_registry, wireguard_manager, firewall_manager,
                         cell_link_manager, _configured_domain, COREFILE_PATH)
        try:
            _wg_addr = wireguard_manager._get_configured_address()
            _wg_subnet = str(ipaddress.ip_network(_wg_addr, strict=False)) if _wg_addr else '10.0.0.0/24'
        except Exception:
            _wg_subnet = '10.0.0.0/24'
        _cell_links = cell_link_manager.list_connections()
        _cell_subnets = [l['vpn_subnet'] for l in _cell_links if l.get('vpn_subnet')]
        data = request.get_json(silent=True) or {}
        existing = peer_registry.get_peer(peer_name)
        if not existing:
            return jsonify({"error": "Peer not found"}), 404

        config_changed = (
            ('internet_access' in data and data['internet_access'] != existing.get('internet_access', True)) or
            ('ip' in data and data['ip'] != existing.get('ip')) or
            ('persistent_keepalive' in data and data['persistent_keepalive'] != existing.get('persistent_keepalive'))
        )

        updates = {k: v for k, v in data.items()}
        if config_changed:
            updates['config_needs_reinstall'] = True

        success = peer_registry.update_peer(peer_name, updates)
        if success:
            updated_peer = peer_registry.get_peer(peer_name)
            if updated_peer:
                firewall_manager.apply_peer_rules(updated_peer['ip'], updated_peer,
                                                  wg_subnet=_wg_subnet, cell_subnets=_cell_subnets)
                firewall_manager.apply_all_dns_rules(peer_registry.list_peers(), COREFILE_PATH, _configured_domain(),
                                                     cell_links=cell_link_manager.list_connections())
            return jsonify({"message": f"Peer {peer_name} updated", "config_changed": config_changed})
        return jsonify({"error": "Update failed"}), 500
    except Exception as e:
        logger.error(f"Error updating peer {peer_name}: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/<peer_name>/route-via', methods=['PUT'])
def set_peer_route_via(peer_name):
    """Route a peer's internet traffic through a connected exit cell.

    Body: {"via_cell": "cellB"} to enable, {"via_cell": null} to disable.
    On enable: updates WG AllowedIPs and adds policy routing rule inside
    cell-wireguard so the peer's packets egress through the exit cell.
    On disable: reverts AllowedIPs and removes the ip rule.
    Also signals the exit cell to add/remove the FORWARD-to-eth0 firewall rule.
    """
    try:
        from app import peer_registry, wireguard_manager, cell_link_manager
        data = request.get_json(silent=True)
        if data is None or 'via_cell' not in data:
            return jsonify({'error': 'via_cell field required (string or null)'}), 400

        via_cell = data['via_cell']
        if via_cell is not None and not isinstance(via_cell, str):
            return jsonify({'error': 'via_cell must be a string or null'}), 400

        peer = peer_registry.get_peer(peer_name)
        if not peer:
            return jsonify({'error': 'Peer not found'}), 404

        peer_ip = peer.get('ip', '').split('/')[0]
        if not peer_ip:
            return jsonify({'error': 'Peer has no IP assigned'}), 400

        old_via = peer.get('route_via')

        # Remove old routing if switching away from a previous exit cell
        if old_via and old_via != via_cell:
            old_link = next(
                (l for l in cell_link_manager.list_connections()
                 if l['cell_name'] == old_via), None
            )
            if old_link:
                wireguard_manager.update_cell_peer_allowed_ips(
                    old_link['public_key'], old_link['vpn_subnet'],
                    add_default_route=False)
            wireguard_manager.remove_peer_route_via(peer_ip)
            try:
                cell_link_manager.set_exit_relay_active(old_via, False)
            except Exception as e:
                logger.warning(f"set_exit_relay_active(False) for {old_via!r} failed: {e}")

        # Apply new routing
        if via_cell:
            link = next(
                (l for l in cell_link_manager.list_connections()
                 if l['cell_name'] == via_cell), None
            )
            if not link:
                return jsonify({'error': f"Cell {via_cell!r} not connected"}), 404
            if link.get('remote_exit_relay_active'):
                return jsonify({
                    'error': (
                        f"Cannot route via '{via_cell}': it is already routing peers "
                        f"through this cell — enabling both directions would create a loop"
                    )
                }), 409
            wireguard_manager.update_cell_peer_allowed_ips(
                link['public_key'], link['vpn_subnet'], add_default_route=True)
            wireguard_manager.apply_peer_route_via(peer_ip, via_wg_ip=link['dns_ip'])
            try:
                cell_link_manager.set_exit_relay_active(via_cell, True)
            except Exception as e:
                logger.warning(f"set_exit_relay_active(True) for {via_cell!r} failed: {e}")

        updated_peer = peer_registry.set_route_via(peer_name, via_cell)
        return jsonify({'message': f"Route-via for '{peer_name}' updated", 'peer': updated_peer})
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        logger.error(f"Error setting route-via for {peer_name}: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/api/peers/<peer_name>/clear-reinstall', methods=['POST'])
def clear_peer_reinstall(peer_name):
    try:
        from app import peer_registry
        peer_registry.clear_reinstall_flag(peer_name)
        return jsonify({"message": "Reinstall flag cleared"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/<peer_name>', methods=['DELETE'])
def remove_peer(peer_name):
    try:
        from app import (peer_registry, wireguard_manager, firewall_manager,
                         email_manager, calendar_manager, file_manager, auth_manager,
                         cell_link_manager, _configured_domain, COREFILE_PATH)
        peer = peer_registry.get_peer(peer_name)
        if not peer:
            return jsonify({"message": f"Peer {peer_name} not found or already removed"})
        peer_ip = peer.get('ip')
        peer_pubkey = peer.get('public_key', '')
        success = peer_registry.remove_peer(peer_name)
        if success:
            if peer_ip:
                firewall_manager.clear_peer_rules(peer_ip)
            firewall_manager.apply_all_dns_rules(peer_registry.list_peers(), COREFILE_PATH, _configured_domain(),
                                                 cell_links=cell_link_manager.list_connections())
            if peer_pubkey:
                try:
                    wireguard_manager.remove_peer(peer_pubkey)
                except Exception as wg_err:
                    logger.warning(f"Peer {peer_name}: WireGuard removal failed (non-fatal): {wg_err}")
        for _cleanup in [
            lambda: email_manager.delete_email_user(peer_name, _configured_domain()),
            lambda: calendar_manager.delete_calendar_user(peer_name),
            lambda: file_manager.delete_user(peer_name),
            lambda: auth_manager.delete_user(peer_name),
        ]:
            try:
                _cleanup()
            except Exception:
                pass
        return jsonify({"message": f"Peer {peer_name} removed successfully"})
    except Exception as e:
        logger.error(f"Error removing peer: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/register', methods=['POST'])
def register_peer():
    try:
        from app import peer_registry
        data = request.get_json(silent=True)
        if data is None:
            return jsonify({"error": "No data provided"}), 400
        return jsonify(peer_registry.register_peer(data))
    except Exception as e:
        logger.error(f"Error registering peer: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/<peer_name>/unregister', methods=['DELETE'])
def unregister_peer(peer_name):
    try:
        from app import peer_registry
        return jsonify(peer_registry.unregister_peer(peer_name))
    except Exception as e:
        logger.error(f"Error unregistering peer: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/peers/<peer_name>/update-ip', methods=['PUT'])
def update_peer_ip_registry(peer_name):
    try:
        from app import peer_registry, routing_manager, wireguard_manager
        data = request.get_json(silent=True)
        new_ip = data.get('ip') if data else None
        if not new_ip:
            return jsonify({"error": "Missing ip"}), 400
        peer = peer_registry.get_peer(peer_name)
        success = peer_registry.update_peer_ip(peer_name, new_ip)
        if success:
            try:
                routing_manager.update_peer_ip(peer_name, new_ip)
            except Exception as e:
                logger.warning(f"RoutingManager update_peer_ip failed: {e}")
            if peer and peer.get('public_key'):
                try:
                    wg_ip = new_ip if '/' in new_ip else f'{new_ip}/32'
                    wireguard_manager.update_peer_ip(peer['public_key'], wg_ip)
                except Exception as e:
                    logger.warning(f"WireGuard update_peer_ip failed: {e}")
            return jsonify({"message": f"IP update received for {peer_name}"})
        return jsonify({"error": f"Peer {peer_name} not found"}), 404
    except Exception as e:
        logger.error(f"Error updating peer IP: {e}")
        return jsonify({"error": str(e)}), 500


@bp.route('/api/ip-update', methods=['POST'])
def ip_update():
    try:
        from app import peer_registry, routing_manager, wireguard_manager
        data = request.get_json(silent=True)
        if data is None:
            return jsonify({"error": "No data provided"}), 400
        peer_name = data.get('peer')
        new_ip = data.get('ip')
        if not peer_name or not new_ip:
            return jsonify({"error": "Missing peer or ip"}), 400
        peer = peer_registry.get_peer(peer_name)
        success = peer_registry.update_peer_ip(peer_name, new_ip)
        if success:
            try:
                routing_manager.update_peer_ip(peer_name, new_ip)
            except Exception as e:
                logger.warning(f"RoutingManager update_peer_ip failed: {e}")
            if peer and peer.get('public_key'):
                try:
                    wg_ip = new_ip if '/' in new_ip else f'{new_ip}/32'
                    wireguard_manager.update_peer_ip(peer['public_key'], wg_ip)
                except Exception as e:
                    logger.warning(f"WireGuard update_peer_ip failed: {e}")
            return jsonify({"message": f"IP update received for {peer_name}"})
        return jsonify({"error": f"Peer {peer_name} not found"}), 404
    except Exception as e:
        logger.error(f"Error handling IP update: {e}")
        return jsonify({"error": str(e)}), 500