pic/api/routes/cells.py

import ipaddress
import logging
import os
from datetime import datetime
from flask import Blueprint, request, jsonify
from cell_link_manager import VALID_SERVICES
logger = logging.getLogger('picell')
bp = Blueprint('cells', __name__)


def _authenticate_peer_cell(req):
    """Return the cell_links record whose vpn_subnet contains the request source IP.

    Source IP is taken from the last X-Forwarded-For entry (appended by Caddy)
    when present, falling back to request.remote_addr.  Also verifies the
    body's from_public_key matches the matched link — defence-in-depth against
    overlapping subnets.
    Returns the matching link dict on success, None on failure.
    """
    from app import cell_link_manager

    candidate = req.remote_addr or ''
    xff = req.headers.get('X-Forwarded-For', '')
    if xff:
        last = xff.split(',')[-1].strip()
        if last:
            candidate = last
    try:
        src_ip = ipaddress.ip_address(candidate.strip())
    except Exception:
        return None

    for link in cell_link_manager.list_connections():
        subnet = link.get('vpn_subnet')
        if not subnet:
            continue
        try:
            if src_ip in ipaddress.ip_network(subnet, strict=False):
                return link
        except Exception:
            continue
    return None

@bp.route('/api/cells/invite', methods=['GET'])
def get_cell_invite():
    try:
        from app import cell_link_manager, config_manager
        identity = config_manager.configs.get('_identity', {})
        cell_name = identity.get('cell_name', os.environ.get('CELL_NAME', 'mycell'))
        domain = identity.get('domain', os.environ.get('CELL_DOMAIN', 'cell'))
        return jsonify(cell_link_manager.generate_invite(cell_name, domain))
    except Exception as e:
        logger.error(f"Error generating cell invite: {e}")
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells/services', methods=['GET'])
def list_shareable_services():
    """Return the list of services that can be shared between cells."""
    try:
        from firewall_manager import SERVICE_IPS
        return jsonify({'services': list(SERVICE_IPS.keys())})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells', methods=['GET'])
def list_cell_connections():
    try:
        from app import cell_link_manager
        return jsonify(cell_link_manager.list_connections())
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells', methods=['POST'])
def add_cell_connection():
    try:
        from app import cell_link_manager
        data = request.get_json(silent=True)
        if not data:
            return jsonify({'error': 'No data provided'}), 400
        for field in ('cell_name', 'public_key', 'vpn_subnet', 'dns_ip', 'domain'):
            if field not in data:
                return jsonify({'error': f'Missing field: {field}'}), 400
        inbound_services = data.get('inbound_services', [])
        link = cell_link_manager.add_connection(data, inbound_services=inbound_services)
        return jsonify({'message': f"Connected to cell '{data['cell_name']}'", 'link': link}), 201
    except ValueError as e:
        return jsonify({'error': str(e)}), 400
    except RuntimeError as e:
        return jsonify({'error': str(e)}), 400
    except Exception as e:
        logger.error(f"Error adding cell connection: {e}")
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells/<cell_name>', methods=['DELETE'])
def remove_cell_connection(cell_name):
    try:
        from app import cell_link_manager
        cell_link_manager.remove_connection(cell_name)
        return jsonify({'message': f"Cell '{cell_name}' disconnected"})
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        logger.error(f"Error removing cell connection: {e}")
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells/<cell_name>/status', methods=['GET'])
def get_cell_connection_status(cell_name):
    try:
        from app import cell_link_manager
        return jsonify(cell_link_manager.get_connection_status(cell_name))
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells/<cell_name>/permissions', methods=['GET'])
def get_cell_permissions(cell_name):
    try:
        from app import cell_link_manager
        perms = cell_link_manager.get_permissions(cell_name)
        return jsonify(perms)
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@bp.route('/api/cells/<cell_name>/permissions', methods=['PUT'])
def update_cell_permissions(cell_name):
    try:
        from app import cell_link_manager, firewall_manager, peer_registry
        from app import COREFILE_PATH
        data = request.get_json(silent=True)
        if not data:
            return jsonify({'error': 'No data provided'}), 400

        # Validate service names in inbound/outbound
        for direction in ('inbound', 'outbound'):
            for service in data.get(direction, {}):
                if service not in VALID_SERVICES:
                    return jsonify({'error': f'Unknown service: {service!r}'}), 400

        inbound = data.get('inbound', {})
        outbound = data.get('outbound', {})
        link = cell_link_manager.update_permissions(cell_name, inbound, outbound)

        # Regenerate Corefile so outbound DNS changes take effect
        try:
            from app import config_manager
            domain = config_manager.configs.get('_identity', {}).get('domain', 'cell')
            peers = peer_registry.list_peers()
            cell_links = cell_link_manager.list_connections()
            firewall_manager.apply_all_dns_rules(peers, COREFILE_PATH, domain,
                                                 cell_links=cell_links)
        except Exception as e:
            logger.warning(f"DNS regen after permission update failed (non-fatal): {e}")

        return jsonify({'message': f"Permissions updated for '{cell_name}'", 'link': link})
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        logger.error(f"Error updating cell permissions: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/api/cells/<cell_name>/exit-offer', methods=['PUT'])
def set_exit_offer(cell_name):
    """Toggle whether this cell offers to route internet for a connected cell.

    Body: {"exit_offered": true|false}
    The new value is persisted and pushed to the remote cell.
    """
    try:
        from app import cell_link_manager
        data = request.get_json(silent=True) or {}
        if 'exit_offered' not in data:
            return jsonify({'error': 'exit_offered field required'}), 400
        link = cell_link_manager.set_exit_offered(cell_name, bool(data['exit_offered']))
        return jsonify({'message': f"Exit offer for '{cell_name}' updated", 'link': link})
    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        logger.error(f"Error setting exit offer: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/api/cells/peer-sync/accept-invite', methods=['POST'])
def peer_sync_accept_invite():
    """Machine-to-machine: a newly-connected cell pushes its own invite for mutual WG pairing.

    Called by Cell A over the LAN (before the WG tunnel exists) immediately after Cell A
    imports Cell B's invite.  Cell B uses this to add Cell A as a WireGuard peer and
    complete the bidirectional tunnel setup without manual admin action on Cell B.

    No session auth — the request arrives before the WG tunnel is up.  Basic sanity
    checks (valid invite format, no subnet/domain conflicts) are applied.  The endpoint
    is idempotent: calling it again for an already-connected cell is a no-op.
    """
    try:
        from app import cell_link_manager
        data = request.get_json(silent=True) or {}
        invite = data.get('invite')
        if not invite or not isinstance(invite, dict):
            return jsonify({'ok': False, 'error': 'invite object required'}), 400

        for field in ('cell_name', 'public_key', 'vpn_subnet', 'dns_ip', 'domain'):
            if field not in invite:
                return jsonify({'ok': False, 'error': f'invite missing field: {field!r}'}), 400

        if invite.get('version') not in (1, None):
            return jsonify({'ok': False, 'error': 'unsupported invite version'}), 400

        link = cell_link_manager.accept_invite(invite)
        return jsonify({'ok': True, 'cell_name': link['cell_name']}), 201
    except ValueError as e:
        return jsonify({'ok': False, 'error': str(e)}), 400
    except RuntimeError as e:
        return jsonify({'ok': False, 'error': str(e)}), 400
    except Exception as e:
        logger.error(f'accept-invite error: {e}')
        return jsonify({'ok': False, 'error': 'internal error'}), 500


@bp.route('/api/cells/peer-sync/permissions', methods=['POST'])
def peer_sync_permissions():
    """Machine-to-machine endpoint: a connected cell pushes its mirrored permission state.

    Auth: source IP must be inside a known cell's vpn_subnet AND the body's
    from_public_key must match that cell's stored public key.
    No session/CSRF required — the WireGuard tunnel is the authentication layer.
    """
    try:
        link = _authenticate_peer_cell(request)
        if not link:
            logger.warning(f"peer-sync: rejected from {request.remote_addr} — no matching cell")
            return jsonify({'ok': False, 'error': 'unauthorized'}), 403

        data = request.get_json(silent=True) or {}

        if data.get('version') != 1:
            return jsonify({'ok': False, 'error': 'unsupported or missing version'}), 400

        sender_pubkey = data.get('from_public_key', '')
        if not sender_pubkey or sender_pubkey != link.get('public_key'):
            logger.warning(
                f"peer-sync: pubkey mismatch from {request.remote_addr} "
                f"(claimed cell={data.get('from_cell')!r})"
            )
            return jsonify({'ok': False, 'error': 'unauthorized'}), 403

        perms = data.get('permissions') or {}
        if not isinstance(perms, dict):
            return jsonify({'ok': False, 'error': 'permissions must be an object'}), 400
        for direction in ('inbound', 'outbound'):
            for svc in (perms.get(direction) or {}):
                if svc not in VALID_SERVICES:
                    return jsonify({'ok': False, 'error': f'unknown service: {svc!r}'}), 400

        exit_offered = bool(data.get('exit_offered', False))
        use_as_exit_relay = bool(data.get('use_as_exit_relay', False))
        from app import cell_link_manager
        cell_link_manager.apply_remote_permissions(sender_pubkey, perms,
                                                   exit_offered=exit_offered,
                                                   use_as_exit_relay=use_as_exit_relay)
        return jsonify({'ok': True, 'applied_at': datetime.utcnow().isoformat()})
    except ValueError as e:
        return jsonify({'ok': False, 'error': str(e)}), 404
    except Exception as e:
        logger.error(f"peer-sync error: {e}")
        return jsonify({'ok': False, 'error': 'internal error'}), 500