3cloud-backend/app/utils/create_workload_container.py

#app/utils/create_workload_container.py
import random
from flask import json, request, jsonify
import requests
from app import db, logger, app
from app.models.models import PortForwarding, Workload, WorkloadHost, WorkloadHostPortMapping, VirtualDataCenter, Label, VolumeWorkloadMapping, ContainerPod, ContainerPodContainer, CloudflareDNSRecord, CloudflareTunnel, WorkloadRequest
import uuid
from app.scheduling_filters import ExcludeAllDisabledHosts, MostAvailableCapacity, ExcludeAllOfflineHosts, DockerCapableHosts, ExcludeHostsWithoutNorthSouthIP, RandomizeHostsOrder
from app.services.cloudflare import CloudflareTunnelManager
from werkzeug.exceptions import abort

from app.utils.standard_responses import api_response

def add_container_workload(request_data):

    # DONT Validate incoming data
    validated_data = request_data.payload
    logger.debug(validated_data)

    request_vdc = VirtualDataCenter.query.filter_by(id=(validated_data['virtual_data_center'])).first()
    if request_vdc is None:
        logger.error(f"Virtual Data Center not found with ID: {validated_data['virtual_data_center']}")
        return api_response(
            success=False,
            status=404,
            message="Virtual Data Center not found",
            error_type="RESOURCE_NOT_FOUND"
        )
    all_workload_hosts = WorkloadHost.query.filter_by(region_id=request_vdc.region.id).all()
    logger.debug(f"Found {len(all_workload_hosts)} hosts eligible for placement")

    filters = [ExcludeHostsWithoutNorthSouthIP(), DockerCapableHosts(), ExcludeAllOfflineHosts(), ExcludeAllDisabledHosts(), MostAvailableCapacity(), RandomizeHostsOrder()]
    filtered_hosts = all_workload_hosts
    for filter in filters:
        filtered_hosts = filter.apply(filtered_hosts)

    if not filtered_hosts:
        logger.error(f"No hosts found in region {request_vdc.region.name} to place this request after filtering")
        return api_response(
            success=False,
            status=400,
            message=f"No hosts found in region {request_vdc.region.name} to place this request",
            error_type="NO_ELIGIBLE_HOSTS"
        )

    selected_host = filtered_hosts[0]
    logger.info(f"Selected host: {selected_host}")

    # --- Step 1: Create a new Pod ---
    new_pod = ContainerPod(
        name=f"Pod_{uuid.uuid4()}",
        workload_host_id=selected_host.id,
        vdc_id=request_vdc.id

    )
    db.session.add(new_pod)
    db.session.flush()  # Flush so we get new_pod.id

    # --- Step 2: Create NSController for the Pod ---
    new_nscontroller_name = f"NSCONTROLLER_{uuid.uuid4()}"
    nscontroller_launch_params = {
        "docker_image": "busybox",
        "container_name": new_nscontroller_name,
        "command": "sleep infinity",
        "networks": [],
        "ports": []
    }
    new_nscontroller = Workload(
        name=new_nscontroller_name,
        workload_type="NSController",
        vdc_id=request_vdc.id,
        launch_params=json.dumps(nscontroller_launch_params),
        workload_host_id=selected_host.id
    )
    db.session.add(new_nscontroller)
    db.session.flush()  # So we have new_nscontroller.id for the pod linkage

    new_pod.nscontroller_workload_id = new_nscontroller.id
    db.session.add(new_pod)
    db.session.flush()

    client_response = []

    # Track Cloudflare configurations for bulk setup
    ingress_mappings = []
    dns_records_to_create = []

    # --- Step 3: Create all Container Workloads inside the Pod ---
    for _container in validated_data['containers']:
        logger.debug(f"Creating container {_container['container_name']}")

        # Build initial container object
        new_container = Workload(
            name=_container['container_name'],
            workload_type="Container",
            vdc_id=request_vdc.id,
            launch_params=json.dumps(_container),
            workload_host_id=selected_host.id
        )
        db.session.add(new_container)
        db.session.flush()  # So we have new_container.id

        # Map container to pod
        mapping = ContainerPodContainer(
            pod_id=new_pod.id,
            container_workload_id=new_container.id
        )
        db.session.add(mapping)

        # --- Step 4: Handle ports and WorkloadHostPortMapping ---
        if 'ports' in _container:
            for port_mapping in _container['ports']:
                internal_port = port_mapping['internal']
                preferred_external_port = port_mapping['external']
                use_dns = port_mapping.get('use_dns', False)

                assigned_external_port = find_available_external_port(
                    selected_host.id, preferred_port=preferred_external_port
                )

                logger.info(f"Assigning internal port:{internal_port} -> external port:{assigned_external_port} for container {new_container.id}")

                # Add this port to NSController launch params
                nscontroller_launch_params['ports'].append({
                    "internal": internal_port,
                    "external": assigned_external_port
                })

                # Save host-level port mapping
                new_port_mapping = WorkloadHostPortMapping(
                    workload_host_id=selected_host.id,
                    container_workload_id=new_container.id,
                    internal_port=internal_port,
                    external_port=assigned_external_port
                )
                db.session.add(new_port_mapping)
        
                # Save the pod-level forwarding rule
                new_forward = PortForwarding(
                    name="f",
                    pod_id=new_pod.id,
                    internal_port=internal_port,
                    external_port=assigned_external_port,
                    protocol=port_mapping.get("protocol", "tcp"),
                    ip_address=selected_host.ip_address_northsouth,
                    workload_id=new_container.id,
                )
                db.session.add(new_forward)
                db.session.flush()  # Ensure we have new_forward.id

                # If use_dns is True for this port, prepare Cloudflare config
                if use_dns:
                    hostname = f"{_container['container_name']}-{internal_port}.hawkvelt.tech"
                    ingress_mappings.append({
                        "dns_hostname": hostname,
                        "local_ip": "127.0.0.1",
                        "local_port": str(internal_port)
                    })
                    dns_records_to_create.append({
                        "hostname": hostname,
                        "forward_id": new_forward.id,
                        "internal_port": internal_port
                    })

        client_response.append({
            "pod_id": str(new_pod.id),
            "container_id": str(new_container.id),
            "container_name": new_container.name,
            "container_status": "pending-allocation",
        })

    # Handle Cloudflare setup if we have any DNS-enabled ports
    if ingress_mappings:
        logger.debug(f"Setting up Cloudflare tunnel with {len(ingress_mappings)} routes")
        
        api_token = "ri6lIjM-aRJBY_xZ82w0Haew93U6YgZYHi5jby1-"
        account_id = "5095a74b62fee53cc5d997c67443bac5"
        zone_id = "e2cafdd8929869d5db885d8824f514b5"
        
        cloudflare_manager = CloudflareTunnelManager(api_token, account_id, zone_id, logger)
        
        # Create a single tunnel with multiple ingress rules
        tunnel_name = f"tun-{new_pod.id}"
        cloudflare_response = cloudflare_manager.setup_tunnel(
            tunnel_name=tunnel_name,
            ingress_mappings=ingress_mappings
        )
        
        logger.debug(cloudflare_response)
        cloudflare_tunnel_token = cloudflare_response['token']
        
        # Create tunnel record in our database
        new_cloudflare_tunnel = CloudflareTunnel(
            token=cloudflare_tunnel_token,
            account_id=account_id,
            tunnel_id=cloudflare_response['tunnel_id'],
            name=cloudflare_response['tunnel_name'],
            tunnel_secret=cloudflare_response['tunnel_secret'],
            nscontroller_workload_id=new_nscontroller.id
        )
        db.session.add(new_cloudflare_tunnel)
        db.session.flush()
        
        logger.debug(dns_records_to_create)
        # Create DNS records for each hostname
        for dns_record in dns_records_to_create:
            # Create DNS record in Cloudflare
            # dns_response = cloudflare_manager.create_dns_record(
            #     hostname=dns_record['hostname'],
            #     target=f"{cloudflare_response['tunnel_id']}.cfargotunnel.com"
            # )
            # TODO - Make sure we are connecting the DNS record in the DB t whjats in Cloudflare so we can delete them later
            # The commented code above is not required becuase the DNS records are created with the tunnel, so we need to query 
            # the setup_tunnel reponse for the DNs info...maybe?
            # Create DNS record in our database
            
            for _createdDNSRecord in cloudflare_response['dns_records_created']:
                logger.debug(_createdDNSRecord)
                if _createdDNSRecord['hostname']==dns_record['hostname']:
                    logger.debug(f"Match {_createdDNSRecord['hostname']}")
                    dns_record_id=_createdDNSRecord['response']['id']

            new_dns_record = CloudflareDNSRecord(
                name="dns",
                zone_id=zone_id,
                dns_record_id=dns_record_id,
                hostname=dns_record['hostname'],
                record_type="CNAME",
                content=f"{cloudflare_response['tunnel_id']}.cfargotunnel.com",
                tunnel_id=new_cloudflare_tunnel.id
            )
            db.session.add(new_dns_record)
            db.session.flush()
            logger.debug(f"Created DNS Record - {hostname}")
            # Link the port forwarding to the DNS record
            port_forward = PortForwarding.query.filter_by(id=dns_record['forward_id']).first()

            if port_forward:
                logger.debug("Found the associated PortForwarding in the DB")
                port_forward.dns_record_id = new_dns_record.id
                db.session.add(port_forward)
                logger.debug(f"Linked DNS record to PortForwarding {dns_record['forward_id']}")

        # Create cloudflared sidecar container
        sidecar_container_name = f"cloudflared-sidecar-{new_pod.id}"
        sidecar_launch_params = {
            "docker_image": "cloudflare/cloudflared:latest",
            "container_name": sidecar_container_name,
            "command": f"tunnel --no-autoupdate run --token {cloudflare_tunnel_token}",
            "network_mode": f"container:{new_nscontroller.name}",
            "restart_policy": "always"
        }

        sidecar_workload = Workload(
            name=sidecar_container_name,
            workload_type="Container",
            vdc_id=request_vdc.id,
            launch_params=json.dumps(sidecar_launch_params),
            workload_host_id=selected_host.id
        )
        db.session.add(sidecar_workload)
        db.session.flush()

        sidecar_mapping = ContainerPodContainer(
            pod_id=new_pod.id,
            container_workload_id=sidecar_workload.id
        )
        db.session.add(sidecar_mapping)

        client_response.append({
            "pod_id": str(new_pod.id),
            "container_id": str(sidecar_workload.id),
            "container_name": sidecar_workload.name,
            "container_status": "pending-allocation",
            "cloudflare_tunnel": {
                "tunnel_id": new_cloudflare_tunnel.tunnel_id,
                "hostnames": [dns['hostname'] for dns in dns_records_to_create]
            }
        })

    # --- Step 5: Update NSController launch params with final ports ---
    new_nscontroller.launch_params = json.dumps(nscontroller_launch_params)
    db.session.add(new_nscontroller)

    # --- Step 6: Commit everything ---
    db.session.commit()

    # --- Step 7: Build and Send Full Pod Payload ---
    pod_payload = build_pod_payload(new_pod)
    headers = {"Content-Type": "application/json"}
    logger.debug(f"Sending full pod payload to websocket server: {pod_payload}")

    websocket_server_response = requests.post(app.config['WEBSOCKET_SERVER_URL'], data=json.dumps(pod_payload), headers=headers)

    if websocket_server_response.status_code == 201:
        logger.info(f"Task created successfully for pod {new_pod.id}")
        for container_info in client_response:
            container = Workload.query.get((container_info['container_id']))
            container.set_status("allocated")
            db.session.add(container)
        new_nscontroller.set_status("allocated")
        db.session.add(new_nscontroller)
        db.session.commit()
    else:
        logger.error(f"Pod creation request failed for pod {new_pod.id}")
        for container_info in client_response:
            container = Workload.query.get((container_info['container_id']))
            container.set_status("failed-allocation")
            db.session.add(container)
        new_nscontroller.set_status("failed-allocation")
        db.session.add(new_nscontroller)
        db.session.commit()


def build_pod_payload(pod):
    nscontroller = Workload.query.get(pod.nscontroller_workload_id)

    container_mappings = ContainerPodContainer.query.filter_by(pod_id=pod.id).all()

    containers = []
    for mapping in container_mappings:
        container = Workload.query.get(mapping.container_workload_id)
        launch_params = json.loads(container.launch_params)

        container_payload = {
            "container_id": str(container.id),
            "docker_image": launch_params.get('docker_image'),
            "container_name": container.name,
            "desired_state": "running"
        }
        # If a command is defined in the launch params, add it
        if 'command' in launch_params and launch_params['command']:
            container_payload['command'] = launch_params['command']
            
        containers.append(container_payload)

    nscontroller_launch_params = json.loads(nscontroller.launch_params)

    pod_payload = {
        "worker_id": str(pod.workload_host_id),
        "task_type": "pod-update",
        "job_details": {
            "pod_id": str(pod.id),
            "nscontroller": {
                "container_id": str(nscontroller.id),
                "docker_image": nscontroller_launch_params.get('docker_image'),
                "command": nscontroller_launch_params.get('command', 'sleep infinity'),
                "networks": nscontroller_launch_params.get('networks', []),
                "ports": nscontroller_launch_params.get('ports', [])
            },
            "containers": containers
        }
    }
    return pod_payload



def find_available_external_port(host_id, preferred_port=None, port_range=(30000, 40000)):
    """Find a free external port on a given host."""
    existing_ports = set(
        mapping.external_port for mapping in WorkloadHostPortMapping.query.filter_by(workload_host_id=host_id).all()
    )

    if preferred_port and preferred_port not in existing_ports:
        return preferred_port

    attempts = 0
    while attempts < 1000:
        port_candidate = random.randint(port_range[0], port_range[1])
        if port_candidate not in existing_ports:
            return port_candidate
        attempts += 1

    raise Exception("No available ports found after 1000 attempts")