gonets/SRV/services/firewalld_service.py

from re import DEBUG
from typing import Dict, List, Optional

from pydbus import SystemBus

DEBUG = True


class FirewalldService:
    def __init__(self):
        self.bus = SystemBus()
        self.proxy = self.bus.get(
            "org.fedoraproject.FirewallD1", "/org/fedoraproject/FirewallD1"
        )

    def get_zones(self) -> List[str]:
        """Get list of all zones."""
        try:
            return self.proxy.getZones()
        except Exception as e:
            raise Exception(f"Failed to get zones: {str(e)}")

    def get_default_zone(self) -> str:
        """Get the default zone."""
        try:
            return self.proxy.getDefaultZone()
        except Exception as e:
            raise Exception(f"Failed to get default zone: {str(e)}")

    def set_default_zone(self, zone: str) -> None:
        """Set the default zone."""
        try:
            self.proxy.setDefaultZone(zone)
        except Exception as e:
            raise Exception(f"Failed to set default zone to {zone}: {str(e)}")

    def add_service(self, zone: str, service: str, timeout: Optional[int] = 0) -> None:
        """Add a service to a zone."""
        try:
            self.proxy.addService(zone, service, timeout)
        except Exception as e:
            raise Exception(f"Failed to add service {service} to zone {zone}: {str(e)}")

    def remove_service(self, zone: str, service: str) -> None:
        """Remove a service from a zone."""
        try:
            self.proxy.removeService(zone, service)
        except Exception as e:
            raise Exception(
                f"Failed to remove service {service} from zone {zone}: {str(e)}"
            )

    def get_services(self, zone: str) -> List[str]:
        """Get list of services in a zone."""
        try:
            return self.proxy.getServices(zone)
        except Exception as e:
            raise Exception(f"Failed to get services for zone {zone}: {str(e)}")

    def add_port(
        self, zone: str, port: int, protocol: str, timeout: Optional[int] = 0
    ) -> None:
        """Add a port to a zone."""
        try:
            if DEBUG:
                print(f"Adding port {port}/{protocol} to zone {zone}")
                print(f"Timeout: {timeout}")
            # Get the zone interface
            zone_proxy = self.bus.get(
                "org.fedoraproject.FirewallD1", "/org/fedoraproject/FirewallD1"
            )
            zone_iface = zone_proxy["org.fedoraproject.FirewallD1.zone"]
            # Convert port to string
            port_str = str(port)
            # Convert timeout to integer
            timeout_int = int(timeout) if timeout is not None else 0
            # Call addPort with the correct signature (sssi)
            zone_iface.addPort(zone, port_str, protocol, timeout_int)
        except Exception as e:
            raise Exception(
                f"Failed to add port {port}/{protocol} to zone {zone}: {str(e)}"
            )

    def remove_port(self, zone: str, port: int, protocol: str) -> None:
        """Remove a port from a zone."""
        try:
            port_str = f"{port}/{protocol}"
            self.proxy.removePort(zone, port_str)
        except Exception as e:
            raise Exception(
                f"Failed to remove port {port}/{protocol} from zone {zone}: {str(e)}"
            )

    def get_ports(self, zone: str) -> List[str]:
        """Get list of ports in a zone."""
        try:
            ports = self.proxy.getPorts(zone)
            # Convert port list items to strings in format "port/protocol"
            return [f"{port[0]}/{port[1]}" for port in ports]
        except Exception as e:
            raise Exception(f"Failed to get ports for zone {zone}: {str(e)}")

    def add_source(self, zone: str, source: str) -> None:
        """Add a source to a zone."""
        try:
            self.proxy.addSource(zone, source)
        except Exception as e:
            raise Exception(f"Failed to add source {source} to zone {zone}: {str(e)}")

    def remove_source(self, zone: str, source: str) -> None:
        """Remove a source from a zone."""
        try:
            self.proxy.removeSource(zone, source)
        except Exception as e:
            raise Exception(
                f"Failed to remove source {source} from zone {zone}: {str(e)}"
            )

    def get_sources(self, zone: str) -> List[str]:
        """Get list of sources in a zone."""
        try:
            return self.proxy.getSources(zone)
        except Exception as e:
            raise Exception(f"Failed to get sources for zone {zone}: {str(e)}")

    def reload(self) -> None:
        """Reload firewall rules."""
        try:
            self.proxy.reload()
        except Exception as e:
            raise Exception(f"Failed to reload firewall rules: {str(e)}")

    def get_zone_settings(self, zone: str) -> Dict:
        """Get all settings for a zone."""
        try:
            return {
                "services": self.get_services(zone),
                "ports": self.get_ports(zone),
                "sources": self.get_sources(zone),
            }
        except Exception as e:
            raise Exception(f"Failed to get zone settings for {zone}: {str(e)}")