#!/usr/bin/env python3
import os
import re
import subprocess
import getpass

import click
import pandas as pd
from console import fg, bg

_ANSIRE = re.compile(r'\x1b\[[0-9;]*m')

ESC = "\033["
RESET = ESC + "0m"
BOLD = ESC + "1m"
DIM = ESC + "2m"
GREEN = ESC + "32m"
RED = ESC + "31m"
YELLOW = ESC + "33m"

SVC_OK = f"{GREEN}●{RESET}"
SVC_INACTIVE = f"{DIM}○{RESET}"
SVC_FAIL = f"{RED}●{RESET}"

SUDOERS_MSG = """
    sudo visudo -f /etc/sudoers.d/ufw-status
youruser ALL=(root) NOPASSWD: /usr/sbin/ufw status, /usr/sbin/ufw status verbose
"""

def strip_ansi(s):
    return _ANSIRE.sub('', s)

def pad_ansi(s, width, left=True):
    vis = strip_ansi(s)
    pad = max(0, width - len(vis))
    return (s + ' ' * pad) if left else (' ' * pad + s)

def run(cmd):
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
    return p.stdout


def original_user():
    try:
        out = run(["logname"]).strip()
        if out:
            return out
    except Exception:
        pass
    # fallback
    return os.environ.get("SUDO_USER") or getpass.getuser()


def check_sudoers_file():
    path = '/etc/sudoers.d/ufw-status'
    if os.path.exists(path):
        return True
    click.echo(f"\n  {RED}✗{RESET} {path} not found\n")
    click.echo(SUDOERS_MSG)
    exit(1)


def parse_ufw():
    """Return dict port->rule_str and list of allowed numeric ports."""
    txt = run(["sudo", "ufw", "status"])
    rules, allowed = {}, set()
    for line in txt.splitlines():
        line = line.strip()
        if not line or re.match(r"^[0-9]+\.", line):
            continue
        parts = re.split(r"\s+", line)
        first = parts[0]
        first_norm = re.sub(r"/(tcp|udp)$", "", first)
        if re.match(r"^\d+$", first_norm):
            rule_str = " ".join(parts[:4]) if len(parts) >= 4 else " ".join(parts)
            rules[first_norm] = rule_str
            if "DENY" not in rule_str and "DENY" not in line:
                allowed.add(int(first_norm))
    allowed_list = sorted(allowed)
    if not rules or not allowed_list:
        click.echo(f"\n  {bg.red}{fg.white} UFW - PROBLEM with FIREWALL {fg.default}{bg.default}")
    return rules, allowed_list


def check_ufw_for_port(rules, port):
    port_s = str(port)
    return rules.get(port_s, "")


def systemctl_status(svc):
    # append .service (systemctl handles either)
    unit = f"{svc}.service"
    out = run(["systemctl", "is-active", unit]).strip()
    return out or "unknown"


def show_warning_if_policy_not_deny():
    txt = run(["ufw", "status", "verbose"])
    if "Default: deny (incoming)" not in txt:
        click.echo(f"\n  {YELLOW}⚠ Check UFW default policies!{RESET}")
        click.echo(txt)


def ss_listening_contains(port, name_substr=None):
    txt = run(["ss", "-tulpn"])
    if str(port) not in txt:
        return False
    if name_substr:
        return name_substr in txt
    return True


@click.command()
@click.option("-n", "--no-decoration", is_flag=True, help="Do not print headers / separators")
def main(no_decoration):
    orig_user = original_user()

    services = [
        ("ufw", 0),
        ("psad", 0),
        ("sshd", 22),
        ("influxdb", 8086),
        ("chrony", 323),
        ("ntp", 123),
        ("mosquitto", 1883),
        (f"syncthing@{orig_user}", 22000),
        ("grafana-server", 3000),
        ("smbd", 445),
        ("elog", 9000),
        ("nginx", 80),
        ("ntfy", 80),
    ]

    check_sudoers_file()
    rules, open_ports = parse_ufw()
    open_ports_left = list(open_ports)

    if not no_decoration:
        click.echo(f"\n  {BOLD}SERVICE{RESET}                 {BOLD}STATUS{RESET}      {BOLD}PORT{RESET}   {BOLD}UFW RULE{RESET}")
        click.echo(f"  {'─' * 60}")

    rows = []
    for svc, port in services:
        status = systemctl_status(svc)
        if status == "active":
            icon, color = SVC_OK, GREEN
        elif status == "inactive":
            icon, color = SVC_INACTIVE, DIM
        else:
            icon, color = SVC_FAIL, RED
        
        ufw_str = ""
        if port != 0:
            ufw_rule = check_ufw_for_port(rules, port)
            if ufw_rule:
                if "DENY" in ufw_rule:
                    ufw_str = f"{RED}{ufw_rule}{RESET}"
                else:
                    ufw_str = f"{GREEN}{ufw_rule}{RESET}"
            else:
                ufw_str = f"{DIM}(not in ufw){RESET}"
            if port in open_ports_left:
                open_ports_left = [p for p in open_ports_left if p != port]
        
        rows.append({
            "service": svc,
            "status": f"{icon} {color}{status}{RESET}",
            "port": "" if port == 0 else str(port),
            "ufw": ufw_str,
        })

    df = pd.DataFrame(rows, columns=["service", "status", "port", "ufw"])

    for _, r in df.iterrows():
        svc = r["service"]
        status_colored = r["status"]
        port = str(r["port"])
        ufw = r["ufw"]

        svc_field = f"  {svc:<20}"
        status_field = pad_ansi(status_colored, 12)
        if port.strip():
            click.echo(f"{svc_field} {status_field}  {port:>5}   {ufw}")
        else:
            click.echo(f"{svc_field} {status_field}  {'—':>5}   {ufw}")

    tlgf_active = False
    ps_txt = run(["ps", "-ef"])
    if f"{orig_user}.conf" in ps_txt and "telegraf" in ps_txt:
        tlgf_active = True
    vadi_active = ss_listening_contains(8200, "telegraf")

    tlgf_icon = SVC_OK if tlgf_active else SVC_INACTIVE
    vadi_icon = SVC_OK if vadi_active else SVC_INACTIVE
    tlgf_str = f"{tlgf_icon} {GREEN}active{RESET}" if tlgf_active else f"{tlgf_icon} {DIM}inactive{RESET}"
    vadi_str = f"{vadi_icon} {GREEN}active{RESET}" if vadi_active else f"{vadi_icon} {DIM}inactive{RESET}"

    tlgf_ufw = check_ufw_for_port(rules, 8200)
    tlgf_ufw_str = f"{GREEN}{tlgf_ufw}{RESET}" if tlgf_ufw else f"{DIM}(not in ufw){RESET}"
    
    click.echo(f"  {'─' * 60}")
    click.echo(f"  {'TELEGRAF@' + orig_user:<20} {tlgf_str:<17}  {'8200':>5}   {tlgf_ufw_str}")
    click.echo(f"  {'VADIM-' + orig_user:<20} {vadi_str:<17}  {'8200':>5}   {tlgf_ufw_str}")

    if not no_decoration:
        click.echo(f"  {'─' * 60}")

    if os.geteuid() == 0:
        if open_ports_left:
            click.echo(f"\n  {YELLOW}⚠ Unmatched open ports: {open_ports_left}{RESET}")
        show_warning_if_policy_not_deny()


if __name__ == "__main__":
    main()
