"""Reverse tunnel manager for device-side tunnel establishment."""
import subprocess
import logging
from typing import List, Dict, Optional
import requests

logger = logging.getLogger(__name__)

class ReverseTunnelManager:
    """
    Manages reverse tunnels on the device.

    Fetches tunnel configuration from backend and establishes
    AWS IoT Secure Tunnels using local proxy.
    """

    def __init__(self, device_dsn: str, api_url: str, cert_path: str, key_path: str):
        self.device_dsn = device_dsn
        self.api_url = api_url
        self.cert_path = cert_path
        self.key_path = key_path
        self.tunnels: List[Dict] = []
        self.tunnel_processes: Dict[int, subprocess.Popen] = {}

    def fetch_tunnel_config(self) -> List[Dict]:
        """Fetch reverse tunnel configuration from backend."""
        try:
            response = requests.get(
                f"{self.api_url}/api/v1/reverse-tunnels/devices/{self.device_dsn}/reverse-tunnels-config",
                cert=(self.cert_path, self.key_path),
                verify=True,
                timeout=30
            )
            response.raise_for_status()
            config = response.json()
            self.tunnels = config['tunnels']
            logger.info(f"Fetched {len(self.tunnels)} tunnel configurations")
            return self.tunnels
        except Exception as e:
            logger.error(f"Failed to fetch tunnel config: {e}")
            return []

    def establish_tunnels(self):
        """Establish all enabled tunnels."""
        for tunnel_config in self.tunnels:
            if not tunnel_config['is_enabled']:
                logger.info(f"Tunnel {tunnel_config['name']} is disabled, skipping")
                continue

            if not tunnel_config.get('tunnel_id') or not tunnel_config.get('destination_access_token'):
                logger.warning(f"Tunnel {tunnel_config['name']} has no AWS tunnel, skipping")
                continue

            self.start_tunnel_agent(tunnel_config)

    def start_tunnel_agent(self, config: Dict) -> Optional[subprocess.Popen]:
        """
        Start AWS IoT Secure Tunnel local proxy agent.

        Downloads and runs localproxy binary if not present.
        """
        try:
            # Ensure localproxy is installed
            self._ensure_localproxy_installed()

            command = [
                'localproxy',
                '-r', 'us-east-1',  # AWS region
                '-s', str(config['local_port']),
                '-t', config['destination_access_token'],
            ]

            logger.info(f"Starting tunnel: {config['name']} -> localhost:{config['local_port']}")

            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            self.tunnel_processes[config['id']] = process
            logger.info(f"Tunnel started (PID: {process.pid}) for {config['public_url']}")
            return process

        except Exception as e:
            logger.error(f"Failed to start tunnel {config['name']}: {e}")
            return None

    def _ensure_localproxy_installed(self):
        """Download localproxy if not present."""
        # Check if localproxy exists
        result = subprocess.run(['which', 'localproxy'], capture_output=True)
        if result.returncode == 0:
            return

        # Download localproxy
        logger.info("Downloading AWS IoT Secure Tunneling localproxy...")
        import platform
        import urllib.request

        arch = platform.machine()
        if arch == 'x86_64':
            url = "https://github.com/aws-samples/aws-iot-securetunneling-localproxy/releases/latest/download/localproxy-linux-x86_64"
        elif arch == 'aarch64':
            url = "https://github.com/aws-samples/aws-iot-securetunneling-localproxy/releases/latest/download/localproxy-linux-arm64"
        else:
            raise Exception(f"Unsupported architecture: {arch}")

        urllib.request.urlretrieve(url, '/usr/local/bin/localproxy')
        os.chmod('/usr/local/bin/localproxy', 0o755)
        logger.info("localproxy installed to /usr/local/bin/localproxy")

    def stop_all_tunnels(self):
        """Stop all running tunnel processes."""
        for tunnel_id, process in self.tunnel_processes.items():
            try:
                process.terminate()
                process.wait(timeout=5)
                logger.info(f"Stopped tunnel {tunnel_id}")
            except Exception as e:
                logger.error(f"Failed to stop tunnel {tunnel_id}: {e}")

from roboticks_sdk.reverse_tunnel_manager import ReverseTunnelManager

class RoboticksDevice:
    def __init__(self, ...):
        # ... existing init
        self.reverse_tunnel_manager = None

    def start(self):
        """Start device services."""
        # ... existing startup code

        # Initialize reverse tunnels
        self._setup_reverse_tunnels()

    def _setup_reverse_tunnels(self):
        """Setup reverse tunnels on device startup."""
        try:
            self.reverse_tunnel_manager = ReverseTunnelManager(
                device_dsn=self.dsn,
                api_url=self.config.api_url,
                cert_path=self.config.cert_path,
                key_path=self.config.key_path
            )

            # Fetch configuration
            tunnels = self.reverse_tunnel_manager.fetch_tunnel_config()

            # Establish tunnels
            if tunnels:
                self.reverse_tunnel_manager.establish_tunnels()
                logger.info(f"Reverse tunnels established: {len(tunnels)}")

        except Exception as e:
            logger.error(f"Failed to setup reverse tunnels: {e}")

    def stop(self):
        """Stop device services."""
        # Stop reverse tunnels
        if self.reverse_tunnel_manager:
            self.reverse_tunnel_manager.stop_all_tunnels()

        # ... existing shutdown code

requests>=2.31.0

import boto3
import json
import base64

iotsecuretunneling = boto3.client('iotsecuretunneling')

def lambda_handler(event, context):
    """
    Proxy WebSocket connections to AWS IoT Secure Tunnels.

    Route: device-{dsn}.fleet.roboticks.io → AWS IoT Tunnel
    """
    # Extract device DSN from hostname
    host = event['headers'].get('host', '')
    dsn = host.split('.')[0].replace('device-', '')

    # Look up tunnel by public_url
    # (requires DynamoDB cache or direct DB query)
    tunnel = get_tunnel_by_dsn(dsn)

    if not tunnel or not tunnel['source_access_token']:
        return {
            'statusCode': 404,
            'body': json.dumps({'error': 'Tunnel not found or not active'})
        }

    # Establish WebSocket connection to AWS IoT Secure Tunnel
    # using source_access_token
    # Forward request through tunnel

    # ... proxy logic ...

    return {
        'statusCode': 200,
        'body': 'Connected'
    }

# proxy_service.py
from fastapi import FastAPI, WebSocket
import asyncio
import boto3

app = FastAPI()

@app.websocket("/")
async def websocket_proxy(websocket: WebSocket):
    """Proxy WebSocket to AWS IoT Secure Tunnel."""
    await websocket.accept()

    # Extract DSN from host header
    host = websocket.headers.get('host')
    dsn = extract_dsn(host)

    # Get tunnel config
    tunnel = await get_tunnel_config(dsn)

    # Connect to AWS IoT Secure Tunnel
    tunnel_ws = await connect_to_iot_tunnel(
        tunnel['tunnel_id'],
        tunnel['source_access_token']
    )

    # Bi-directional proxy
    await asyncio.gather(
        forward_client_to_tunnel(websocket, tunnel_ws),
        forward_tunnel_to_client(tunnel_ws, websocket)
    )

docker build -t reverse-tunnel-proxy .
docker push ECR_URL/reverse-tunnel-proxy:latest

# Deploy to ECS or K8s

from celery.schedules import crontab

beat_schedule = {
    'cleanup-expired-reverse-tunnels': {
        'task': 'cleanup_expired_reverse_tunnels',
        'schedule': crontab(minute=0),  # Every hour at :00
    },
    'rotate-expiring-reverse-tunnels': {
        'task': 'rotate_expiring_reverse_tunnels',
        'schedule': crontab(minute='*/30'),  # Every 30 minutes
    },
}

celery -A app.celery_app beat --loglevel=info