#!/usr/bin/env python3

import os
import signal
import argparse
import logging
import re
import time
from threading import Thread
from wsgiref.simple_server import make_server
from prometheus_client import make_wsgi_app, Counter, Gauge

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# Aggregated, labeled metric (safe: labels have small domain)
auth_sessions = Counter(
    'auth_sessions_total',
    'Total auth sessions started',
    ['source', 'method']
)

# Optional, user-labeled metric. Disabled by default because usernames are high-cardinality.
auth_sessions_by_user = Counter(
    'auth_sessions_by_user_total',
    'Total auth sessions by username (high-cardinality — enable with --export-usernames)',
    ['source', 'method', 'username']
)

# exporter health/observability
parse_errors = Counter('auth_exporter_parse_errors_total', 'Total parse errors')
uptime = Gauge('auth_exporter_uptime_seconds', 'Exporter uptime in seconds')
user_label_dropped = Counter('auth_exporter_user_labels_dropped_total', 'Number of username labels dropped due to max-user-labels limit')

# Current active sessions (gauge) — safe labels with small domain
active_sessions = Gauge('auth_sessions_active', 'Current active auth sessions', ['source'])
# Optional per-user active gauge (high-cardinality — enabled with --export-usernames)
active_sessions_by_user = Gauge('auth_sessions_active_by_user', 'Current active auth sessions by username (high-cardinality — enable with --export-usernames)', ['source', 'username'])

# more tolerant username pattern (allows -, ., @, digits)
USERNAME = r'[\w\-\.\@]+'
IP = r'[\d\.]+|[\da-fA-F:]+'

pam_ssh_session = re.compile(
    rf'sshd\[\d+\]: Accepted (password|publickey) for ({USERNAME}) from ({IP}) port \d+ ssh2'
)
tty_pattern = re.compile(
    rf'login\[\d+\]: pam_unix\(login:session\): session opened for user ({USERNAME}) by \(uid=\d+\)'
)
systemd_logind_session = re.compile(
    rf'systemd-logind\[\d+\]: New session \d+ of user ({USERNAME})\.'
)
pam_systemd_session = re.compile(
    rf'systemd: pam_unix\(systemd-user:session\): session opened for user ({USERNAME})(?:\(uid=\d+\))? by \(uid=\d+\)'
)
pam_sudo_session = re.compile(
    rf'sudo: pam_unix\(sudo[^)]*:session\): session opened for user ({USERNAME})(?:\(uid=\d+\))? by (?:({USERNAME})|uid=\d+|\(uid=\d+\))'
)

# Generic 'session closed' pattern (captures service and username when pam_unix logs a session close)
pam_session_closed = re.compile(rf'pam_unix\((?P<service>[^:]+):session\): session closed for user ({USERNAME})')

# Track seen usernames to limit cardinality if username export is enabled
_seen_usernames = set()

class TailReader:
    def __init__(self, path, start_at_end=True):
        self.path = path
        self.file = None
        self.inode = None
        self.start_at_end = start_at_end
        self.open_file()

    def open_file(self):
        try:
            self.file = open(self.path, 'r', errors='ignore')
            stat = os.fstat(self.file.fileno())
            self.inode = (stat.st_ino, stat.st_dev)
            if self.start_at_end:
                self.file.seek(0, os.SEEK_END)
            logging.info("Opened %s (inode=%s)", self.path, self.inode)
        except Exception as e:
            logging.error("Failed to open %s: %s", self.path, e)
            raise

    def check_rotate(self):
        try:
            stat = os.stat(self.path)
            current = (stat.st_ino, stat.st_dev)
            if current != self.inode:
                logging.info("Detected rotation/truncate for %s (old=%s new=%s)", self.path, self.inode, current)
                self.file.close()
                self.open_file()
            else:
                # handle copytruncate: file same inode but possibly truncated
                try:
                    pos = self.file.tell()
                    if pos > stat.st_size:
                        logging.info("Detected file truncated for %s (pos=%d size=%d), seeking to %d", self.path, pos, stat.st_size, stat.st_size)
                        self.file.seek(stat.st_size, os.SEEK_SET)
                except Exception as e:
                    logging.error("Error checking for truncation: %s", e)
        except FileNotFoundError:
            # file temporarily missing (rotation); try reopen later
            logging.warning("File %s not found while monitoring; will retry", self.path)
            time.sleep(1)
        except Exception as e:
            logging.error("Error checking rotation: %s", e)

    def follow_once(self):
        # yield any new lines
        while True:
            line = self.file.readline()
            if not line:
                break
            yield line

def parse_line(line, export_usernames=False, max_user_labels=100):
    try:
        # First handle session closed lines so we can decrement active gauges
        m_close = pam_session_closed.search(line)
        if m_close:
            service = m_close.group('service')
            username = m_close.group(2)
            # map service to a source label
            if service.startswith('sshd') or service == 'sshd':
                source = 'ssh'
            elif service.startswith('login') or service == 'login':
                source = 'tty'
            elif service.startswith('systemd') or service.startswith('systemd-user'):
                source = 'systemd-user'
            elif service.startswith('sudo'):
                source = 'sudo'
            else:
                source = None

            if source:
                # decrement the aggregated active gauge (but avoid negative values when possible)
                try:
                    child = active_sessions.labels(source=source)
                    val = child._value.get()
                    if val and val > 0:
                        child.dec()
                except Exception:
                    # fallback: attempt to decrement anyway
                    try:
                        active_sessions.labels(source=source).dec()
                    except Exception:
                        pass

                if export_usernames and username in _seen_usernames:
                    try:
                        childu = active_sessions_by_user.labels(source=source, username=username)
                        v = childu._value.get()
                        if v and v > 0:
                            childu.dec()
                    except Exception:
                        try:
                            active_sessions_by_user.labels(source=source, username=username).dec()
                        except Exception:
                            pass
            # nothing else to do for closed lines
            return 'closed'

        # ssh (captures method and username)
        m = pam_ssh_session.search(line)
        if m:
            method = m.group(1)  # 'password' or 'publickey'
            username = m.group(2)
            auth_sessions.labels(source='ssh', method=method).inc()
            # increment active sessions gauge (aggregated by source)
            try:
                active_sessions.labels(source='ssh').inc()
            except Exception:
                pass
            if export_usernames:
                if username not in _seen_usernames:
                    if len(_seen_usernames) >= max_user_labels:
                        user_label_dropped.inc()
                    else:
                        _seen_usernames.add(username)
                if username in _seen_usernames:
                    auth_sessions_by_user.labels(source='ssh', method=method, username=username).inc()
                    try:
                        active_sessions_by_user.labels(source='ssh', username=username).inc()
                    except Exception:
                        pass
            return 'ssh'

        # tty
        m = tty_pattern.search(line)
        if m:
            username = m.group(1)
            auth_sessions.labels(source='tty', method='local').inc()
            if export_usernames:
                if username not in _seen_usernames:
                    if len(_seen_usernames) >= max_user_labels:
                        user_label_dropped.inc()
                    else:
                        _seen_usernames.add(username)
                if username in _seen_usernames:
                    auth_sessions_by_user.labels(source='tty', method='local', username=username).inc()
            return 'tty'

        # systemd-logind
        m = systemd_logind_session.search(line)
        if m:
            username = m.group(1)
            auth_sessions.labels(source='systemd-logind', method='pam').inc()
            if export_usernames:
                if username not in _seen_usernames:
                    if len(_seen_usernames) >= max_user_labels:
                        user_label_dropped.inc()
                    else:
                        _seen_usernames.add(username)
                if username in _seen_usernames:
                    auth_sessions_by_user.labels(source='systemd-logind', method='pam', username=username).inc()
            return 'systemd-logind'

        # systemd user session via pam
        m = pam_systemd_session.search(line)
        if m:
            username = m.group(1)
            auth_sessions.labels(source='systemd-user', method='pam').inc()
            if export_usernames:
                if username not in _seen_usernames:
                    if len(_seen_usernames) >= max_user_labels:
                        user_label_dropped.inc()
                    else:
                        _seen_usernames.add(username)
                if username in _seen_usernames:
                    auth_sessions_by_user.labels(source='systemd-user', method='pam', username=username).inc()
            return 'systemd-user'

        # sudo
        m = pam_sudo_session.search(line)
        if m:
            # sudo regex captures the target username (group 1) or similar; normalize to 'sudo'
            username = m.group(1)
            auth_sessions.labels(source='sudo', method='sudo').inc()
            if export_usernames:
                if username not in _seen_usernames:
                    if len(_seen_usernames) >= max_user_labels:
                        user_label_dropped.inc()
                    else:
                        _seen_usernames.add(username)
                if username in _seen_usernames:
                    auth_sessions_by_user.labels(source='sudo', method='sudo', username=username).inc()
            return 'sudo'

    except Exception:
        parse_errors.inc()
        logging.exception("Error parsing line")
    return None

running = True
start_time = time.time()
httpd = None
_server_thread = None

def handle_signal(signum, frame):
    global running, httpd
    logging.info("Received signal %s, shutting down", signum)
    running = False
    try:
        if httpd:
            logging.info("Shutting down HTTP server")
            httpd.shutdown()
    except Exception:
        logging.exception("Error shutting down HTTP server")

def main():
    global running
    parser = argparse.ArgumentParser(description="Auth log prometheus exporter")
    parser.add_argument('--logfile', default='/var/log/auth.log', help='Path to auth log')
    parser.add_argument('--port', type=int, default=9100, help='Prometheus metrics port')
    parser.add_argument('--start-at-beginning', action='store_true', help='Parse file from beginning on first run')
    parser.add_argument('--interval', type=float, default=1.0, help='Poll interval seconds')
    parser.add_argument('--export-usernames', action='store_true', help='Enable exporting username labels (DANGEROUS: may increase cardinality)')
    parser.add_argument('--max-user-labels', type=int, default=100, help='Maximum distinct username labels to export when --export-usernames is enabled')
    args = parser.parse_args()

    # NOTE: reading /var/log/auth.log usually requires root
    if not os.access(args.logfile, os.R_OK):
        logging.warning("No read access to %s. Exporter may need to run as root or use systemd-journal.", args.logfile)

    # Start Prometheus WSGI app in a daemon thread so it does not block process exit.
    global httpd, _server_thread
    app = make_wsgi_app()
    httpd = make_server('', args.port, app)
    _server_thread = Thread(target=httpd.serve_forever, name="prometheus-http")
    _server_thread.daemon = True
    _server_thread.start()
    logging.info("Starting Prometheus session exporter on port %d...", args.port)

    tail = TailReader(args.logfile, start_at_end=not args.start_at_beginning)

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    try:
        while running:
            tail.check_rotate()
            for line in tail.follow_once():
                parse_line(line, export_usernames=args.export_usernames, max_user_labels=args.max_user_labels)
            uptime.set(time.time() - start_time)
            time.sleep(args.interval)
    except Exception:
        logging.exception("Unhandled error in main loop")
    finally:
        logging.info("Shutting down exporter, cleaning up resources")
        try:
            if tail.file:
                tail.file.close()
        except Exception:
            logging.exception("Error closing tail file")
        try:
            if httpd:
                httpd.shutdown()
        except Exception:
            logging.exception("Error shutting down HTTP server")
        try:
            if _server_thread:
                _server_thread.join(timeout=2)
        except Exception:
            logging.exception("Error joining HTTP server thread")
        logging.info("Exporter stopped")

if __name__ == '__main__':
    main()