#!/root/.local/bin/uv run
#
#session optional pam_exec.so seteuid /usr/local/bin/login-notify.py
#
# chmod +x login-notify.py
# EDIT: /etc/pam.d/sshd  AND APPEND:
# session optional pam_exec.so seteuid /path/to/login-notify.py
#
# /// script
# requires-python = ">=3.10"
# dependencies = ["click", "requests"]
# ///

DEBUG_LOG = "/tmp/login-notify-debug.log"


def _debug(msg: str):
    try:
        import datetime
        with open(DEBUG_LOG, "a") as f:
            f.write(f"{datetime.datetime.now()}: {msg}\n")
    except Exception:
        pass


_debug("imports starting")

import os
import socket
import subprocess as sp
import datetime as dt
import configparser
import shutil
from pathlib import Path

import click
import requests

_debug("imports done")

DEFAULT_CONFIG_PATH = Path("~/.config/influxdb/totalconfig.conf").expanduser()
DEFAULT_NTFY_SERVER = "https://ntfy.sh"


def _send_ntfy(
    topic: str,
    message: str,
    title: str | None = None,
    priority: int = 4,
    server: str | None = None,
    username: str | None = None,
    password: str | None = None,
) -> bool:
    server = server or DEFAULT_NTFY_SERVER
    if topic.startswith("http"):
        url = topic
    else:
        url = f"{server.rstrip('/')}/{topic}"
    headers = {"Priority": str(priority)}
    if title:
        headers["Title"] = title
    auth = (username, password) if username and password else None
    response = requests.post(url, data=message.encode('utf-8'), headers=headers, auth=auth, timeout=10)
    response.raise_for_status()
    return True


def _load_config() -> dict[str, str | None]:
    config: dict[str, str | None] = {'server': None, 'username': None, 'password': None, 'topic': None}
    if not DEFAULT_CONFIG_PATH.exists():
        return config
    parser = configparser.ConfigParser()
    parser.read(DEFAULT_CONFIG_PATH)
    if 'ntfy login' in parser:
        section = parser['ntfy login']
        config['server'] = section.get('server')
        config['username'] = section.get('username')
        config['password'] = section.get('password')
        config['topic'] = section.get('topic')
    return config


def _get_log_file() -> str:
    user = os.environ.get('PAM_USER', os.environ.get('USER', 'unknown'))
    return f"/tmp/login-notify-{user}.log"


def _log(msg: str):
    try:
        with open(_get_log_file(), "a") as f:
            f.write(f"{msg}\n")
    except Exception as e:
        _debug(f"log failed: {e}")


def _get_mac(ip: str) -> str:
    if ip.startswith("127.") or ip == "::1":
        return "localhost"
    try:
        cmd = f"arping -c 1 -w 1 -r {ip}"
        _log(f"   {cmd} with uid={os.getuid()}")
        res = sp.check_output(cmd.split())
        mac = res.decode("utf8").strip()
        _log(f"   mac: {mac}")
        return mac
    except Exception as e:
        _log(f"   mac lookup failed: {e}")
        return "?"


PAM_FILE = "/etc/pam.d/sshd"
INSTALL_PATH = "/usr/local/bin/login-notify.py"
PAM_LINE = "session optional pam_exec.so seteuid /usr/local/bin/login-notify.py"


@click.group(invoke_without_command=True)
@click.pass_context
def cli(ctx):
    """Login notification via ntfy.

    Script location: /usr/local/bin/login-notify.py
    PAM config: /etc/pam.d/sshd
      session optional pam_exec.so seteuid /usr/local/bin/login-notify.py

    Setup commands:
      sudo ntfy access user topic_x rw
      ntfy pub topic_x "test"
    """
    if ctx.invoked_subcommand is None:
        ctx.invoke(run)


@cli.command()
@click.option('--ntfy-topic', default=None, help='ntfy topic (default: from config or hostname)')
@click.option('--dry-run', is_flag=True, help='Print message without sending')
def run(ntfy_topic, dry_run):
    """Send login notification via ntfy."""
    _log("NEWLOG-----------------------------------")
    _log("     start")

    host = socket.gethostname()
    rhost = os.environ.get('PAM_RHOST', 'unknown')
    user = os.environ.get('PAM_USER', 'unknown')
    pam_type = os.environ.get('PAM_TYPE', 'unknown')
    tty = os.environ.get('PAM_TTY', '')

    _log(f"     host={host}, rhost={rhost}, user={user}, type={pam_type}")

    mac = _get_mac(rhost) if rhost != 'unknown' else '?'
    time_str = dt.datetime.now().strftime("%a %H:%M:%S")

    parts = [f"type:{pam_type}", f"mac:{mac}"]
    if tty:
        parts.append(f"tty:{tty}")
    parts.append(time_str)
    message = f"from {rhost}\n" + " | ".join(parts)
    title = f"{user}@{host}"

    if dry_run:
        click.echo(f"Dry run - would send:\nTitle: {title}\n{message}")
        return

    config = _load_config()
    topic = ntfy_topic or config.get('topic') or host
    server = config.get('server')
    username = config.get('username')
    password = config.get('password')

    try:
        _send_ntfy(topic, message, title=title, server=server, username=username, password=password)
        click.echo(f"Notification sent to {topic}")
    except Exception as e:
        click.echo(f"Failed to send notification: {e}", err=True)

    _log(f"------------------{dt.datetime.now()}--")
    _log(f"{user}@{host} < {rhost} / type={pam_type} / tty={tty}")


@cli.command()
def install():
    """Install login-notify.py to system."""
    script_path = Path(__file__).resolve()
    script_stat = script_path.stat()
    script_size = script_stat.st_size
    script_mtime = dt.datetime.fromtimestamp(script_stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S")

    click.echo("Step 1: Checking PAM configuration file...")
    if not Path(PAM_FILE).exists():
        click.echo(f"  ERROR: {PAM_FILE} not found", err=True)
        return
    click.echo(f"  OK: {PAM_FILE} exists")

    click.echo("\nStep 2: Checking PAM configuration line...")
    pam_content = Path(PAM_FILE).read_text()
    if PAM_LINE in pam_content:
        click.echo(f"  OK: PAM line already present")
    else:
        click.echo(f"  WARNING: PAM line not found")
        click.echo(f"  To add manually, append to {PAM_FILE}:")
        click.echo(f"    {PAM_LINE}")

    click.echo("\nStep 3: Installing script...")
    click.echo(f"  Current script: {script_size} bytes, {script_mtime}")
    dest = Path(INSTALL_PATH)
    if dest.exists():
        stat = dest.stat()
        mtime = dt.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S")
        size = stat.st_size
        click.echo(f"  Installed file: {size} bytes, {mtime}")
        if not click.confirm("  Overwrite?", default=False):
            click.echo("  Skipped.")
            return
    else:
        click.echo(f"  Target: {INSTALL_PATH} (new)")

    try:
        shutil.copy2(script_path, dest)
        dest.chmod(0o755)
        click.echo(f"  OK: Installed to {INSTALL_PATH}")
    except PermissionError:
        click.echo(f"  ERROR: Permission denied. Run with sudo.", err=True)
        return

    click.echo("\nStep 4: Config file (optional)...")
    if DEFAULT_CONFIG_PATH.exists():
        click.echo(f"  OK: {DEFAULT_CONFIG_PATH} exists")
    else:
        click.echo(f"  Create {DEFAULT_CONFIG_PATH} with:")
        click.echo('    [ntfy login]')
        click.echo('    server = https://ntfy.sh')
        click.echo('    username = your_user')
        click.echo('    password = your_password')
        click.echo('    topic = ssh_logins')

    click.echo("\nInstallation complete.")


if __name__ == "__main__":
    cli()
