"""
Firewalld 自定义服务管理工具
用于创建和修改 /etc/firewalld/services/ 目录下的自定义服务 XML 文件
"""
import os
from pathlib import Path
from xml.dom import minidom
import xml.etree.ElementTree as ET
from sunhpc.logs import Logs

slog = Logs(name="Firewd", show_time=True)

class FWManager:
    """Firewalld 服务文件管理器"""
    
    # firewalld 服务目录路径
    SYSTEM_SERVICES_DIR = Path('/usr/lib/firewalld/services')
    CUSTOM_SERVICES_DIR = Path('/etc/firewalld/services')
    
    def __init__(self, service_name):
        """
        初始化服务管理器
        
        Args:
            service_name: 服务名称（不含 .xml 后缀）
        """
        self.service_name = service_name
        self.service_file = self.CUSTOM_SERVICES_DIR / f"{service_name}.xml"
        self.xml_tree = None
        self.root = None
        
    def _ensure_custom_dir(self):
        """确保自定义服务目录存在"""
        self.CUSTOM_SERVICES_DIR.mkdir(parents=True, exist_ok=True)
        
    def _create_xml_structure(self):
        """创建基本的 XML 结构"""
        # 创建根元素
        self.root = ET.Element('service')
        
        # 添加 short 元素
        short = ET.SubElement(self.root, 'short')
        short.text = self.service_name.capitalize()
        
        # 添加 description 元素
        description = ET.SubElement(self.root, 'description')
        description.text = f"Custom service: {self.service_name}"
        
        # 添加 port 元素（稍后添加）
        # 添加 protocol 元素（稍后添加）
        
        self.xml_tree = ET.ElementTree(self.root)
        
    def _prettify_xml(self, elem):
        """
        格式化 XML 输出
        
        Args:
            elem: XML 元素
            
        Returns:
            str: 格式化后的 XML 字符串
        """
        rough_string = ET.tostring(elem, 'utf-8')
        reparsed = minidom.parseString(rough_string)
        return reparsed.toprettyxml(indent='  ')
        
    def _save_xml(self):
        """保存 XML 文件到自定义服务目录"""
        self._ensure_custom_dir()
        
        # 格式化 XML
        xml_str = self._prettify_xml(self.root)
        
        # 写入文件
        with open(self.service_file, 'w', encoding='utf-8') as f:
            # 移除多余的 XML 声明（firewalld 不需要）
            if xml_str.startswith('<?xml'):
                xml_str = xml_str.split('\n', 1)[1]
            f.write(xml_str)
            
        # 设置正确的权限（644）
        os.chmod(self.service_file, 0o644)
        
    def load(self):
        """
        加载现有服务文件或创建新文件
        
        Returns:
            bool: 是否成功加载/创建
        """
        if self.service_file.exists():
            # 加载现有文件
            try:
                self.xml_tree = ET.parse(self.service_file)
                self.root = self.xml_tree.getroot()
                slog.info(f"✓ 已加载现有服务: {self.service_file}")
                return True
            except ET.ParseError as e:
                slog.error(f"✗ 解析 XML 文件失败: {e}")
                return False
        else:
            # 创建新文件
            slog.info(f"服务不存在，将创建新服务: {self.service_name}")
            self._create_xml_structure()
            return True
            
    def set_short_description(self, short_text):
        """
        设置服务简短描述
        
        Args:
            short_text: 简短描述文本
        """
        short_elem = self.root.find('short')
        if short_elem is None:
            short_elem = ET.SubElement(self.root, 'short')
        short_elem.text = short_text
        slog.info(f"✓ 设置简短描述: {short_text}")
        
    def set_description(self, desc_text):
        """
        设置服务详细描述
        
        Args:
            desc_text: 详细描述文本
        """
        desc_elem = self.root.find('description')
        if desc_elem is None:
            desc_elem = ET.SubElement(self.root, 'description')
        desc_elem.text = desc_text
        slog.info(f"✓ 设置详细描述: {desc_text}")
        
    def add_port(self, port, protocol='tcp'):
        """
        添加端口规则
        
        Args:
            port: 端口号或端口范围（例如：'8080' 或 '3000-3010'）
            protocol: 协议类型（tcp/udp/sctp/dccp）
        """
        # 查找是否已存在相同的端口规则
        for port_elem in self.root.findall('port'):
            if port_elem.get('port') == str(port) and port_elem.get('protocol') == protocol:
                slog.info(f"⚠ 跳过,已存在, 端口 {port}/{protocol}")
                return
                
        # 添加新端口
        port_elem = ET.SubElement(self.root, 'port')
        port_elem.set('port', str(port))
        port_elem.set('protocol', protocol)
        slog.info(f"✓ 添加端口: {port}/{protocol}")
        
    def remove_port(self, port, protocol='tcp'):
        """
        移除端口规则
        
        Args:
            port: 端口号或端口范围
            protocol: 协议类型
        """
        for port_elem in self.root.findall('port'):
            if port_elem.get('port') == str(port) and port_elem.get('protocol') == protocol:
                self.root.remove(port_elem)
                slog.info(f"✓ 移除端口: {port}/{protocol}")
                return
        slog.info(f"⚠ 未找到端口: {port}/{protocol}")
        
    def add_protocol(self, protocol_name):
        """
        添加协议规则（IP 协议，不是传输层协议）
        
        Args:
            protocol_name: 协议名称（如 ipv6-icmp, icmp, igmp 等）
        """
        for proto_elem in self.root.findall('protocol'):
            if proto_elem.get('name') == protocol_name:
                slog.info(f"⚠ 协议 {protocol_name} 已存在，跳过")
                return
                
        proto_elem = ET.SubElement(self.root, 'protocol')
        proto_elem.set('name', protocol_name)
        slog.info(f"✓ 添加协议: {protocol_name}")
        
    def remove_protocol(self, protocol_name):
        """
        移除协议规则
        
        Args:
            protocol_name: 协议名称
        """
        for proto_elem in self.root.findall('protocol'):
            if proto_elem.get('name') == protocol_name:
                self.root.remove(proto_elem)
                slog.info(f"✓ 移除协议: {protocol_name}")
                return
        slog.info(f"⚠ 未找到协议: {protocol_name}")
        
    def add_module(self, module_name):
        """
        添加内核模块依赖
        
        Args:
            module_name: 模块名称（如 nf_conntrack_ftp）
        """
        for module_elem in self.root.findall('module'):
            if module_elem.get('name') == module_name:
                slog.info(f"⚠ 模块 {module_name} 已存在，跳过")
                return
                
        module_elem = ET.SubElement(self.root, 'module')
        module_elem.set('name', module_name)
        slog.info(f"✓ 添加模块: {module_name}")
        
    def remove_module(self, module_name):
        """
        移除内核模块依赖
        
        Args:
            module_name: 模块名称
        """
        for module_elem in self.root.findall('module'):
            if module_elem.get('name') == module_name:
                self.root.remove(module_elem)
                slog.info(f"✓ 移除模块: {module_name}")
                return
        slog.info(f"⚠ 未找到模块: {module_name}")
        
    def get_ports(self):
        """获取所有端口配置"""
        ports = []
        for port_elem in self.root.findall('port'):
            ports.append({
                'port': port_elem.get('port'),
                'protocol': port_elem.get('protocol')
            })
        return ports
        
    def get_protocols(self):
        """获取所有协议配置"""
        protocols = []
        for proto_elem in self.root.findall('protocol'):
            protocols.append(proto_elem.get('name'))
        return protocols
        
    def display_info(self):
        """显示服务信息"""
        slog.info("\n" + "="*50)
        slog.info(f"服务名称: {self.service_name}")
        slog.info(f"配置文件: {self.service_file}")
        
        short = self.root.find('short')
        if short is not None and short.text:
            slog.info(f"简短描述: {short.text}")
            
        desc = self.root.find('description')
        if desc is not None and desc.text:
            slog.info(f"详细描述: {desc.text}")
            
        ports = self.get_ports()
        if ports:
            slog.info("\n端口规则:")
            for port in ports:
                slog.info(f"  - {port['port']}/{port['protocol']}")
                
        protocols = self.get_protocols()
        if protocols:
            print("\n协议规则:")
            for proto in protocols:
                print(f"  - {proto}")
                
                modules = [m.get('name') for m in self.root.findall('module')]
        if modules:
            slog.info("\n内核模块:")
            for module in modules:
                slog.info(f"  - {module}")
        slog.info("="*50 + "\n")
        
    def save(self):
        """保存修改"""
        self._save_xml()

def main():
    """示例用法"""
    # 创建或修改名为 'myapp' 的服务
    manager = FWManager('myapp')
    
    # 加载或创建服务文件
    if not manager.load():
        return
        
    # 设置描述信息
    manager.set_short_description('My Application Service')
    manager.set_description('Custom service for my web application with API and admin ports')
    
    # 添加端口
    manager.add_port(8080, 'tcp')      # Web 服务端口
    manager.add_port(8443, 'tcp')      # HTTPS 替代端口
    manager.add_port(9090, 'tcp')      # API 端口
    manager.add_port(3000, 'tcp')      # 管理界面端口
    manager.add_port(5000-5010, 'tcp') # 端口范围示例
    manager.add_port(9000, 'udp')      # UDP 端口
    
    # 添加协议（如果需要）
    # manager.add_protocol('icmp')
    
    # 添加内核模块（如果需要，例如 FTP 辅助模块）
    # manager.add_module('nf_conntrack_ftp')
    
    # 显示当前配置
    manager.display_info()
    
    # 保存修改
    manager.save()
    
    # 示例：移除端口
    # manager.remove_port(5000-5010, 'tcp')
    # manager.save()
    
    slog.info("\n✅ 操作完成！")
    slog.info("提示：")
    slog.info("1. 运行 'sudo firewall-cmd --reload' 重新加载防火墙")
    slog.info("2. 使用 'sudo firewall-cmd --get-services' 查看所有服务")
    slog.info("3. 使用 'sudo firewall-cmd --info-service=myapp' 查看服务详情")
    slog.info(f"4. 使用 'sudo firewall-cmd --add-service={manager.service_name}' 启用服务")

if __name__ == '__main__':
    main()
