"""
System / OS-level tools.
- Firewall (iptables / firewalld / ufw)
- Swap setup
- Time sync (chrony)
- SSH hardening
- Hostname & hosts
- Disk / fstab
- Limits (ulimit)
- Sysctl tuning
"""
from . import register, Generator, Field
from . import bash_header, quote, bool_str, yes


# ========================== Firewall ==========================
class Firewall(Generator):
    id = "firewall"
    title = "防火墙 (iptables / firewalld / ufw)"
    category = "system"
    icon = "🧱"
    tags = ["firewall", "iptables", "firewalld", "ufw"]
    description = "根据发行版自动选择 firewalld (RHEL) 或 ufw (Debian),开放指定端口。"
    warnings = ["修改防火墙可能断连,建议在带外有 console 的情况下执行。"]
    fields = [
        Field("ports", "要开放的端口 (TCP)", "text", default="22,80,443",
              help="逗号或空格分隔,如 22 80 443 或 22,80,443。"),
        Field("udp_ports", "UDP 端口", "text", default=""),
        Field("trusted_ip", "完全信任的 IP/CIDR", "text", default="",
              help="填 192.168.1.0/24 可对该网段放行所有端口。"),
        Field("ssh_port", "SSH 端口", "number", default="22", min_=1, max_=65535),
        Field("deny_incoming", "默认拒绝入站", "checkbox", default="yes"),
        Field("allow_outgoing", "允许所有出站", "checkbox", default="yes"),
    ]

    def render(self, p):
        ports = p.get("ports", "22,80,443").replace(",", " ").split()
        udps = p.get("udp_ports", "").replace(",", " ").split() if p.get("udp_ports") else []
        trusted = p.get("trusted_ip", "").strip()
        ssh = str(p.get("ssh_port", "22"))
        deny_in = bool_str(p.get("deny_incoming", True))
        allow_out = bool_str(p.get("allow_outgoing", True))
        ports_str = " ".join([quote(p_) for p_ in ports])
        udps_str = " ".join([quote(p_) for p_ in udps]) if udps else ""
        out = [bash_header(self.title)]
        out.append('log "Configuring firewall..."')
        # Build the case/esac block as a list of lines then join — avoids the
        # "implicit string concatenation" gotcha when mixing `+` with bare strings.
        fw = []
        fw.append('case "$DISTRO_ID" in')
        fw.append('    ubuntu|debian)')
        fw.append('        PKG_INSTALL ufw')
        fw.append('        ufw --force reset')
        fw.append('        ufw default ' + ('deny' if deny_in else 'allow') + ' incoming')
        fw.append('        ufw default ' + ('allow' if allow_out else 'deny') + ' outgoing')
        for p_ in ports:
            if p_ != ssh:
                fw.append(f'        ufw allow {p_}/tcp')
        for p_ in udps:
            fw.append(f'        ufw allow {p_}/udp')
        if trusted:
            fw.append(f'        ufw allow from {trusted} to any port 0:65535')
        fw.append('        ufw allow ' + ssh + '/tcp comment "ssh"')
        fw.append('        ufw --force enable')
        fw.append('        ;;')
        fw.append('    centos|rhel|rocky|almalinux|ol|fedora)')
        fw.append('        PKG_INSTALL firewalld')
        fw.append('        systemctl enable --now firewalld')
        for p_ in ports:
            fw.append(f'        firewall-cmd --permanent --add-port={p_}/tcp')
        for p_ in udps:
            fw.append(f'        firewall-cmd --permanent --add-port={p_}/udp')
        if trusted:
            fw.append(f'        firewall-cmd --permanent --add-source={trusted}')
            fw.append('        firewall-cmd --permanent --add-port=0-65535/tcp')
            fw.append('        firewall-cmd --permanent --add-port=0-65535/udp')
        fw.append('        firewall-cmd --reload')
        fw.append('        ;;')
        fw.append('    *)')
        fw.append('        warn "Auto-firewall not supported on $DISTRO_ID, falling back to iptables"')
        fw.append('        PKG_INSTALL iptables-persistent || PKG_INSTALL iptables-services')
        fw.append('        iptables -P INPUT ' + ('DROP' if deny_in else 'ACCEPT'))
        fw.append('        iptables -A INPUT -i lo -j ACCEPT')
        fw.append('        iptables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT')
        for p_ in ports:
            fw.append(f'        iptables -A INPUT -p tcp --dport {p_} -j ACCEPT')
        if trusted:
            fw.append(f'        iptables -A INPUT -s {trusted} -j ACCEPT')
        fw.append('        netfilter-persistent save || service iptables save')
        fw.append('        ;;')
        fw.append('esac')
        out.append('\n'.join(fw))
        out.append('log "Done."')
        return "\n".join(out) + "\n"


# ========================== Swap ==========================
class Swap(Generator):
    id = "swap"
    title = "Swap 交换分区/文件"
    category = "system"
    icon = "💾"
    tags = ["swap"]
    description = "创建 swap 文件并启用。"
    fields = [
        Field("size_gb", "大小 (GB)", "number", default="2", min_=1, max_=128),
        Field("path", "Swap 文件路径", "text", default="/swapfile"),
        Field("swappiness", "Swappiness", "number", default="10", min_=0, max_=100),
    ]

    def render(self, p):
        sz = str(p.get("size_gb", "2"))
        path = p.get("path", "/swapfile")
        sw = str(p.get("swappiness", "10"))
        out = [bash_header(self.title)]
        out.append('log "Creating ' + sz + 'GB swap at ' + path + '..."')
        out.append('if [ -f "' + path + '" ]; then swapoff ' + path + ' && rm -f ' + path + '; fi')
        out.append('fallocate -l ' + sz + 'G ' + path + ' || dd if=/dev/zero of=' + path + ' bs=1M count=$(((' + sz + '*1024))) status=none')
        out.append('chmod 600 ' + path)
        out.append('mkswap ' + path)
        out.append('swapon ' + path)
        out.append('grep -q "' + path + ' " /etc/fstab || echo "' + path + ' none swap sw 0 0" >> /etc/fstab')
        out.append('sysctl -w vm.swappiness=' + sw)
        out.append('grep -q "vm.swappiness" /etc/sysctl.conf || echo "vm.swappiness=' + sw + '" >> /etc/sysctl.conf')
        out.append('swapon --show')
        return "\n".join(out) + "\n"


# ========================== Time sync (chrony) ==========================
class Chrony(Generator):
    id = "chrony"
    title = "时间同步 (chrony)"
    category = "system"
    icon = "🕐"
    tags = ["time", "ntp", "chrony"]
    description = "安装并配置 chrony,使用国内 NTP 源。"
    fields = [
        Field("servers", "NTP 服务器", "textarea",
              default="ntp.aliyun.com\nntp1.aliyun.com\nntp2.aliyun.com\ntime.cloudflare.com",
              help="一行一个。"),
        Field("makestep_threshold", "makestep 阈值 (秒)", "text", default="1.0 3"),
        Field("allow_subnet", "允许同步的子网", "text", default="192.168.0.0/16",
              help="仅本机同步: 留空 = 127.0.0.0/8"),
    ]

    def render(self, p):
        srvs = [s.strip() for s in p.get("servers", "").splitlines() if s.strip()]
        ms = p.get("makestep_threshold", "1.0 3")
        allow = p.get("allow_subnet", "192.168.0.0/16").strip()
        out = [bash_header(self.title)]
        out.append('log "Installing chrony..."')
        out.append('PKG_INSTALL chrony')
        out.append('cp /etc/chrony/chrony.conf /etc/chrony/chrony.conf.bak.$(date +%s) || true')
        out.append('cat > /etc/chrony/chrony.conf <<CFG_EOF\n'
                   'pool ntp.ubuntu.org        maxsources 4\n'
                   'pool 0.pool.ntp.org         maxsources 1\n'
                   'pool 1.pool.ntp.org         maxsources 1\n'
                   'pool 2.pool.ntp.org         maxsources 1\n'
                   'makestep ' + ms + '\n'
                   'rtcsync\n'
                   'logdir /var/log/chrony\n'
                   + ('allow ' + allow + '\n' if allow else 'allow 127.0.0.1\n') +
                   'CFG_EOF')
        for s in srvs:
            out.append('echo "pool ' + s + ' iburst maxsources 4" >> /etc/chrony/chrony.conf')
        out.append('systemctl enable --now chrony')
        out.append('sleep 2')
        out.append('chronyc tracking')
        out.append('chronyc sources -v')
        return "\n".join(out) + "\n"


# ========================== SSH Hardening ==========================
class SSHHardening(Generator):
    id = "ssh-hardening"
    title = "SSH 服务加固"
    category = "security"
    icon = "🔐"
    tags = ["ssh", "security", "hardening"]
    description = "修改 sshd_config:禁 root 登录、改端口、密钥认证、登录横幅。"
    warnings = [
        "修改前请确保有可登录的用户和密钥,否则会锁死!",
        "建议保留当前 SSH 会话,新开一个窗口测试登录成功后再关闭。",
    ]
    fields = [
        Field("port", "新 SSH 端口", "number", default="2222", min_=1, max_=65535),
        Field("permit_root_login", "允许 root 登录", "select", default="no",
              options=["yes", "no", "prohibit-password", "without-password", "forced-commands-only"]),
        Field("password_auth", "允许密码认证", "checkbox", default="no"),
        Field("pubkey_auth", "公钥认证", "checkbox", default="yes"),
        Field("allow_users", "允许登录的用户 (逗号分隔)", "text", default=""),
        Field("allow_groups", "允许登录的组", "text", default=""),
        Field("max_auth_tries", "最大认证尝试", "number", default="3"),
        Field("max_sessions", "最大会话数", "number", default="5"),
        Field("client_alive_interval", "心跳间隔", "number", default="300"),
        Field("client_alive_count_max", "心跳失败次数", "number", default="2"),
        Field("banner_text", "登录横幅", "textarea", default="Authorised access only. All activity is logged."),
    ]

    def render(self, p):
        port = str(p.get("port", "2222"))
        proot = p.get("permit_root_login", "no")
        passw = bool_str(p.get("password_auth"))
        pubkey = bool_str(p.get("pubkey_auth", True))
        allow_users = p.get("allow_users", "").strip()
        allow_groups = p.get("allow_groups", "").strip()
        mt = str(p.get("max_auth_tries", "3"))
        ms = str(p.get("max_sessions", "5"))
        cai = str(p.get("client_alive_interval", "300"))
        cacm = str(p.get("client_alive_count_max", "2"))
        banner = p.get("banner_text", "").strip()
        out = [bash_header(self.title)]
        out.append('log "Backing up /etc/ssh/sshd_config..."')
        out.append('cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak.$(date +%s)')
        out.append('log "Patching sshd_config..."')
        sshd_lines = [
            'Port ' + port,
            'PermitRootLogin ' + proot,
            'PasswordAuthentication ' + ('yes' if passw else 'no'),
            'PubkeyAuthentication ' + ('yes' if pubkey else 'no'),
            'MaxAuthTries ' + mt,
            'MaxSessions ' + ms,
            'ClientAliveInterval ' + cai,
            'ClientAliveCountMax ' + cacm,
            'X11Forwarding no',
            'AllowTcpForwarding no',
            'PermitEmptyPasswords no',
            'UsePAM yes',
        ]
        if allow_users:
            sshd_lines.append('AllowUsers ' + allow_users)
        if allow_groups:
            sshd_lines.append('AllowGroups ' + allow_groups)
        out.append('cat > /etc/ssh/sshd_config.d/00-shellgen.conf <<CFG_EOF\n' + '\n'.join(sshd_lines) + '\nCFG_EOF')
        if banner:
            out.append('cat > /etc/ssh/banner <<BANNER_EOF\n' + banner + '\nBANNER_EOF')
            out.append('echo "Banner /etc/ssh/banner" >> /etc/ssh/sshd_config.d/00-shellgen.conf')
        out.append('sshd -t')
        out.append('systemctl reload sshd || systemctl restart sshd')
        out.append('log "Now listening on :' + port + ' — keep your session open and test from another window!"')
        return "\n".join(out) + "\n"


# ========================== Sysctl tuning ==========================
class SysctlTuning(Generator):
    id = "sysctl-tuning"
    title = "Sysctl 内核参数调优"
    category = "system"
    icon = "🧬"
    tags = ["sysctl", "kernel", "tuning"]
    description = "常用 sysctl 参数:BBR 拥塞控制、TCP 缓冲区、文件句柄、IPv4 转发。"
    fields = [
        Field("enable_bbr", "启用 BBR", "checkbox", default="yes"),
        Field("tcp_tw_reuse", "TCP TIME_WAIT 复用", "checkbox", default="yes"),
        Field("tcp_fastopen", "TCP Fast Open", "checkbox", default="yes"),
        Field("ip_forward", "IPv4 转发", "checkbox", default="no"),
        Field("net_core_somaxconn", "SOMAXCONN", "number", default="4096"),
        Field("file_max", "fs.file-max", "number", default="2097152"),
        Field("max_map_count", "vm.max_map_count (ES 需 262144)", "number", default="262144"),
        Field("apply_immediately", "立即 sysctl -p", "checkbox", default="yes"),
    ]

    def render(self, p):
        bbr = bool_str(p.get("enable_bbr", True))
        tw = bool_str(p.get("tcp_tw_reuse", True))
        tfo = bool_str(p.get("tcp_fastopen", True))
        fwd = bool_str(p.get("ip_forward"))
        somaxconn = str(p.get("net_core_somaxconn", "4096"))
        fmax = str(p.get("file_max", "2097152"))
        mmc = str(p.get("max_map_count", "262144"))
        apply = bool_str(p.get("apply_immediately", True))
        out = [bash_header(self.title)]
        out.append('log "Writing /etc/sysctl.d/99-shellgen.conf..."')
        out.append('cat > /etc/sysctl.d/99-shellgen.conf <<SYSCTL_EOF\n'
                   'fs.file-max = ' + fmax + '\n'
                   'vm.max_map_count = ' + mmc + '\n'
                   'net.core.somaxconn = ' + somaxconn + '\n'
                   'net.core.netdev_max_backlog = 16384\n'
                   'net.ipv4.tcp_max_syn_backlog = 8192\n'
                   'net.ipv4.tcp_slow_start_after_idle = 0\n'
                   'net.ipv4.tcp_tw_reuse = ' + ('1' if tw else '0') + '\n'
                   'net.ipv4.tcp_fin_timeout = 15\n'
                   'net.ipv4.tcp_keepalive_time = 600\n'
                   'net.ipv4.tcp_keepalive_intvl = 30\n'
                   'net.ipv4.tcp_keepalive_probes = 3\n'
                   'net.ipv4.tcp_fastopen = ' + ('3' if tfo else '0') + '\n'
                   'net.ipv4.ip_forward = ' + ('1' if fwd else '0') + '\n'
                   'net.ipv4.conf.all.rp_filter = 1\n'
                   'net.ipv4.conf.default.rp_filter = 1\n'
                   'net.ipv4.conf.all.accept_source_route = 0\n'
                   'net.ipv4.conf.default.accept_source_route = 0\n'
                   'net.ipv4.icmp_echo_ignore_broadcasts = 1\n'
                   'net.ipv4.conf.all.send_redirects = 0\n'
                   'net.ipv4.conf.all.accept_redirects = 0\n'
                   'net.ipv4.conf.all.secure_redirects = 0\n'
                   'net.ipv6.conf.all.accept_redirects = 0\n'
                   'SYSCTL_EOF')
        if bbr:
            out.append('cat >> /etc/sysctl.d/99-shellgen.conf <<SYSCTL_EOF\n'
                       'net.core.default_qdisc = fq\n'
                       'net.ipv4.tcp_congestion_control = bbr\n'
                       'SYSCTL_EOF')
        if apply:
            out.append('sysctl -p /etc/sysctl.d/99-shellgen.conf')
        return "\n".join(out) + "\n"


# ========================== ulimit / limits.conf ==========================
class Ulimit(Generator):
    id = "ulimit"
    title = "limits.conf 资源限制"
    category = "system"
    icon = "📊"
    tags = ["ulimit", "limits"]
    description = "为指定用户/组写 /etc/security/limits.d/*.conf。"
    fields = [
        Field("target", "对象 (user 或 @group)", "text", default="@app",
              help="@ 表示组;不带前缀为用户。"),
        Field("nofile", "nofile (打开文件数)", "number", default="65536"),
        Field("nproc", "nproc (进程数)", "number", default="65536"),
        Field("memlock", "memlock (KB, 无限 = unlimited)", "text", default="unlimited"),
        Field("cpu", "cpu (分钟, 无限 = unlimited)", "text", default="unlimited"),
        Field("data", "data (KB)", "text", default="unlimited"),
    ]

    def render(self, p):
        t = p.get("target", "@app")
        nofile = str(p.get("nofile", "65536"))
        nproc = str(p.get("nproc", "65536"))
        memlock = p.get("memlock", "unlimited")
        cpu = p.get("cpu", "unlimited")
        data = p.get("data", "unlimited")
        out = [bash_header(self.title)]
        out.append('cat > /etc/security/limits.d/99-shellgen.conf <<LIM_EOF\n'
                   + t + ' soft nofile ' + nofile + '\n'
                   + t + ' hard nofile ' + nofile + '\n'
                   + t + ' soft nproc ' + nproc + '\n'
                   + t + ' hard nproc ' + nproc + '\n'
                   + t + ' soft memlock ' + memlock + '\n'
                   + t + ' hard memlock ' + memlock + '\n'
                   + t + ' soft cpu ' + cpu + '\n'
                   + t + ' hard cpu ' + cpu + '\n'
                   + t + ' soft data ' + data + '\n'
                   + t + ' hard data ' + data + '\n'
                   'LIM_EOF')
        out.append('log "Re-login required for new limits to take effect."')
        return "\n".join(out) + "\n"


# ========================== Hostname / hosts ==========================
class Hostname(Generator):
    id = "hostname"
    title = "Hostname 与 /etc/hosts"
    category = "system"
    icon = "🏷️"
    tags = ["hostname", "hosts"]
    description = "修改主机名并更新 /etc/hosts。"
    fields = [
        Field("hostname", "新主机名", "text", default="server01",
              help="FQDN 格式: server01.example.com, 短名会自动取 . 前部分。"),
        Field("extra_hosts", "额外 hosts 行", "textarea", default="",
              help="每行: <ip> <hostname> [alias...]"),
    ]

    def render(self, p):
        h = p.get("hostname", "server01")
        extras = [e for e in p.get("extra_hosts", "").splitlines() if e.strip()]
        out = [bash_header(self.title)]
        out.append('hostnamectl set-hostname ' + h)
        out.append('cp /etc/hosts /etc/hosts.bak.$(date +%s)')
        out.append('# Ensure 127.0.0.1 contains this hostname')
        out.append('grep -q "' + h + '" /etc/hosts || sed -i "s/^127.0.0.1\\(\\s\\+\\)localhost.*/& ' + h + '/" /etc/hosts')
        for e in extras:
            out.append('echo "' + e + '" >> /etc/hosts')
        out.append('hostname -f')
        return "\n".join(out) + "\n"


# ========================== Disk auto-mount via fstab ==========================
class AutoMount(Generator):
    id = "automount"
    title = "磁盘自动挂载 (fstab)"
    category = "system"
    icon = "💽"
    tags = ["mount", "fstab", "disk"]
    description = "为指定磁盘创建 fstab 条目并挂载。"
    fields = [
        Field("device", "设备 (UUID=/dev/sda1 或 LABEL=...)", "text", default="UUID=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
              help="可用 lsblk -f 或 blkid 查询。"),
        Field("mount_point", "挂载点", "text", default="/data"),
        Field("fstype", "文件系统", "select", default="ext4",
              options=["ext4", "ext3", "xfs", "btrfs", "ntfs", "vfat", "nfs", "cifs"]),
        Field("mount_options", "挂载选项", "text", default="defaults,nofail"),
        Field("dump", "dump 字段", "number", default="0"),
        Field("pass", "fsck 顺序", "number", default="2"),
    ]

    def render(self, p):
        dev = p.get("device", "UUID=xxx")
        mp = p.get("mount_point", "/data")
        fs = p.get("fstype", "ext4")
        opts = p.get("mount_options", "defaults,nofail")
        dump = str(p.get("dump", "0"))
        pss = str(p.get("pass", "2"))
        out = [bash_header(self.title)]
        out.append('mkdir -p ' + mp)
        out.append('echo "' + dev + '  ' + mp + '  ' + fs + '  ' + opts + '  ' + dump + '  ' + pss + '" >> /etc/fstab')
        out.append('mount -a')
        out.append('df -h ' + mp)
        return "\n".join(out) + "\n"


for _g in [Firewall, Swap, Chrony, SSHHardening, SysctlTuning, Ulimit, Hostname, AutoMount]:
    register(_g())