firegex-traffic-viewer/backend/routers/setup.py

from fastapi import APIRouter, HTTPException, UploadFile, File
from pydantic import BaseModel
import json
from typing import List, Optional
from utils.models import StatusMessageModel
from routers import nfproxy, nfregex, porthijack, firewall

class ServiceConfig(BaseModel):
    name: str
    port: int
    proto: str
    ip_int: str
    fail_open: bool = True

class PortHijackServiceConfig(BaseModel):
    name: str
    public_port: int
    proxy_port: int
    proto: str
    ip_src: str
    ip_dst: str

class FirewallRuleConfig(BaseModel):
    mode: str
    src: str
    dst: str
    in_int: str
    out_int: str
    proto: str
    sport: str
    dport: str

class SetupConfig(BaseModel):
    services: Optional[List[ServiceConfig]] = []
    porthijack: Optional[List[PortHijackServiceConfig]] = []
    firewall: Optional[List[FirewallRuleConfig]] = []

class SetupResponse(BaseModel):
    status: str
    services_created: int = 0
    porthijack_created: int = 0
    firewall_created: int = 0
    errors: List[str] = []

app = APIRouter()

@app.post("/import", response_model=SetupResponse)
async def import_setup(config: SetupConfig):
    """
    Import services and rules from a setup configuration.
    Creates basic services without filters or regex rules.
    """
    errors = []
    services_count = 0
    porthijack_count = 0
    firewall_count = 0

    # Import Services
    if config.services:
        for service_config in config.services:
            try:
                # Determine which module to use based on protocol
                # HTTP -> NFProxy, TCP/UDP -> can use either (prefer NFProxy)
                if service_config.proto in ["tcp", "http", "udp"]:
                    # Create NFProxy service
                    try:
                        add_form = nfproxy.ServiceAddForm(
                            name=service_config.name,
                            port=service_config.port,
                            proto=service_config.proto,
                            ip_int=service_config.ip_int,
                            fail_open=service_config.fail_open
                        )
                        result = await nfproxy.add_service(add_form)
                        
                        if result.status == "ok":
                            services_count += 1
                        else:
                            errors.append(f"Service '{service_config.name}': Failed to create")
                    except Exception as e:
                        errors.append(f"Service '{service_config.name}': {str(e)}")
                else:
                    errors.append(f"Service '{service_config.name}': Unsupported protocol '{service_config.proto}'")
                    
            except Exception as e:
                errors.append(f"Service '{service_config.name}': {str(e)}")

    # Import PortHijack services
    if config.porthijack:
        for service_config in config.porthijack:
            try:
                add_form = porthijack.ServiceAddForm(
                    name=service_config.name,
                    public_port=service_config.public_port,
                    proxy_port=service_config.proxy_port,
                    proto=service_config.proto,
                    ip_src=service_config.ip_src,
                    ip_dst=service_config.ip_dst
                )
                result = await porthijack.add_service(add_form)
                
                if result.status == "ok":
                    porthijack_count += 1
                else:
                    errors.append(f"PortHijack service '{service_config.name}': Failed to create")
            except Exception as e:
                errors.append(f"PortHijack service '{service_config.name}': {str(e)}")

    # Import Firewall rules
    if config.firewall:
        for rule_config in config.firewall:
            try:
                rule_form = firewall.RuleFormAdd(
                    mode=rule_config.mode,
                    src=rule_config.src,
                    dst=rule_config.dst,
                    in_int=rule_config.in_int,
                    out_int=rule_config.out_int,
                    proto=rule_config.proto,
                    sport=rule_config.sport,
                    dport=rule_config.dport
                )
                await firewall.add_rule(rule_form)
                firewall_count += 1
            except Exception as e:
                errors.append(f"Firewall rule: {str(e)}")

    return SetupResponse(
        status="ok" if len(errors) == 0 else "partial",
        services_created=services_count,
        porthijack_created=porthijack_count,
        firewall_created=firewall_count,
        errors=errors
    )

@app.post("/import/file")
async def import_setup_file(file: UploadFile = File(...)):
    """
    Import services from an uploaded JSON file.
    """
    try:
        content = await file.read()
        config_dict = json.loads(content.decode('utf-8'))
        config = SetupConfig(**config_dict)
        return await import_setup(config)
    except json.JSONDecodeError as e:
        raise HTTPException(status_code=400, detail=f"Invalid JSON: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error processing file: {str(e)}")

@app.get("/export")
async def export_setup():
    """
    Export all current services and rules as a JSON configuration.
    Exports only service definitions without filters or regexes.
    """
    config = {
        "services": [],
        "porthijack": [],
        "firewall": []
    }

    # Export NFProxy services
    try:
        nfproxy_services = await nfproxy.get_services()
        for service in nfproxy_services:
            config["services"].append({
                "name": service.name,
                "port": service.port,
                "proto": service.proto,
                "ip_int": service.ip_int,
                "fail_open": service.fail_open
            })
    except:
        pass

    # Export NFRegex services
    try:
        nfregex_services = await nfregex.get_services()
        for service in nfregex_services:
            config["services"].append({
                "name": service.name,
                "port": service.port,
                "proto": service.proto,
                "ip_int": service.ip_int,
                "fail_open": service.fail_open
            })
    except:
        pass

    # Export PortHijack services
    try:
        porthijack_services = await porthijack.get_services()
        for service in porthijack_services:
            config["porthijack"].append({
                "name": service.name,
                "public_port": service.public_port,
                "proxy_port": service.proxy_port,
                "proto": service.proto,
                "ip_src": service.ip_src,
                "ip_dst": service.ip_dst
            })
    except:
        pass

    # Export Firewall rules
    try:
        fw_rules = await firewall.get_rules()
        for rule in fw_rules:
            config["firewall"].append({
                "mode": rule.mode,
                "src": rule.src,
                "dst": rule.dst,
                "in_int": rule.in_int,
                "out_int": rule.out_int,
                "proto": rule.proto,
                "sport": rule.sport,
                "dport": rule.dport
            })
    except:
        pass

    return config