#!/usr/bin/env python3
# Emits Prometheus textfile metrics describing pending apt updates.

import os
import re
import subprocess
import sys
import tempfile
import time

OUT_DIR = '/var/lib/prometheus/node-exporter'
OUT_PATH = os.path.join(OUT_DIR, 'apt_updates.prom')

# A pending update is treated as kernel-related if its package name starts
# with one of these prefixes. Covers both meta-packages (linux-image-generic,
# linux-image-amd64, ...) and specific images (linux-image-5.15.0-91-generic).
KERNEL_PREFIXES = ('linux-image-',)

# Regex for `apt list --upgradable` lines. Example lines:
#   foo/jammy-updates,jammy-security 1.2.3-1ubuntu2 amd64 [upgradable from: 1.2.0-1ubuntu1]
#   libbar:i386/jammy 5.0.1 i386 [upgradable from: 4.9.8]
LIST_LINE_RE = re.compile(
    r'^(?P<name>\S+?)/\S+\s+(?P<new>\S+)\s+\S+\s+\[upgradable from:\s+(?P<old>\S+)\]'
)

# Specific-version kernel image package: linux-image-<release>, where <release>
# starts with a digit (skips meta-packages like linux-image-generic).
INSTALLED_KERNEL_PKG_RE = re.compile(r'^linux-image-(\d.*)$')


def strip_arch(name):
    # apt list shows multi-arch packages as `name:arch`. Drop the suffix.
    return name.split(':', 1)[0]


def upstream(version):
    # Strip Debian epoch (N:) and revision (-Nfoo) so we compare upstream only.
    if ':' in version:
        version = version.split(':', 1)[1]
    if '-' in version:
        version = version.rsplit('-', 1)[0]
    return version


def numeric_components(v):
    # Leading numeric components, stopping at the first non-numeric segment.
    # e.g. '7.0.5+ds' -> [7, 0, 5]; '1.2.3~rc1' -> [1, 2, 3]; 'foo' -> [].
    out = []
    for part in v.split('.'):
        m = re.match(r'^(\d+)', part)
        if not m:
            break
        out.append(int(m.group(1)))
    return out


def classify(old, new):
    # Anything that fails to parse cleanly is flagged as 'major' (safer
    # to over-flag a tiny change than under-flag a breaking one).
    o, n = upstream(old), upstream(new)
    if o == n:
        return 'patch'  # epoch/revision-only change
    oc, nc = numeric_components(o), numeric_components(n)
    if not oc or not nc:
        return 'major'
    width = max(len(oc), len(nc), 3)
    oc += [0] * (width - len(oc))
    nc += [0] * (width - len(nc))
    if oc[0] != nc[0]:
        return 'major'
    if oc[1] != nc[1]:
        return 'minor'
    if oc[2] != nc[2]:
        return 'patch'
    # First three numeric components match but upstream strings differ
    # (e.g. '1.2.3+git123' vs '1.2.3+git456'). Bucket as patch.
    return 'patch'


def esc(v):
    # Escape characters that are special in Prometheus label values.
    return v.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')


def installed_kernel_releases():
    # All installed specific-version kernel images, keyed by release string
    # (the part after 'linux-image-', e.g. '5.15.0-91-generic') -> dpkg version.
    # Use a tab separator so the Status field (which contains spaces, e.g.
    # 'install ok installed' vs 'deinstall ok config-files') stays intact.
    r = subprocess.run(
        ['dpkg-query', '-W', '-f=${Package}\t${Version}\t${Status}\n', 'linux-image-*'],
        capture_output=True, text=True, check=False,
    )
    out = {}
    if r.returncode != 0:
        return out
    for line in r.stdout.splitlines():
        parts = line.split('\t')
        if len(parts) != 3:
            continue
        pkg, ver, status = parts
        if status != 'install ok installed':
            continue
        m = INSTALLED_KERNEL_PKG_RE.match(pkg)
        if not m:
            continue
        out[m.group(1)] = ver
    return out


def release_sort_key(release):
    # Turn '5.15.0-91-generic' into ((5, 15, 0, 91), 'generic') for ordering.
    head, _, flavor = release.partition('-generic')
    if flavor == '' and head != release:
        # Matches '<numeric>-generic' exactly.
        flavor = 'generic'
        nums = head
    else:
        # Fallback: split on the last dash, treat tail as flavor.
        nums, _, flavor = release.rpartition('-')
        if not nums:
            nums, flavor = release, ''
    parts = []
    for seg in re.split(r'[.\-]', nums):
        m = re.match(r'^(\d+)', seg)
        parts.append(int(m.group(1)) if m else 0)
    return (tuple(parts), flavor)


def reboot_required(running, kernels):
    # Reboot needed if any installed kernel image sorts higher than the
    # currently running one. Mirrors the Arch script's intent: ask the
    # package manager what's installed and compare to uname -r.
    if not kernels:
        return 0
    running_key = release_sort_key(running)
    for release in kernels:
        if release_sort_key(release) > running_key:
            return 1
    return 0


def main():
    r = subprocess.run(
        ['apt', 'list', '--upgradable'],
        capture_output=True, text=True, check=False,
        env={**os.environ, 'LC_ALL': 'C', 'LANG': 'C'},
    )
    if r.returncode != 0:
        sys.exit(f'apt list --upgradable failed: {r.stderr.strip()}')

    counts = {'patch': 0, 'minor': 0, 'major': 0}
    kernel = 0
    entries = []

    for line in r.stdout.splitlines():
        m = LIST_LINE_RE.match(line)
        if not m:
            continue
        name = strip_arch(m.group('name'))
        old = m.group('old')
        new = m.group('new')
        sev = classify(old, new)
        counts[sev] += 1
        entries.append((name, old, new, sev))
        if name.startswith(KERNEL_PREFIXES):
            kernel = 1

    total = sum(counts.values())
    ts = int(time.time())

    info_lines = [
        '# HELP apt_pending_update_info One series per pending package; details are in the labels.',
        '# TYPE apt_pending_update_info gauge',
    ]
    for name, old, new, sev in entries:
        info_lines.append(
            f'apt_pending_update_info{{'
            f'package="{esc(name)}",'
            f'old_version="{esc(old)}",'
            f'new_version="{esc(new)}",'
            f'severity="{sev}"'
            f'}} 1'
        )
    info_block = '\n'.join(info_lines) + '\n'

    kernels = installed_kernel_releases()
    kernel_info_lines = [
        '# HELP apt_kernel_installed_info Installed kernel images, one series per release.',
        '# TYPE apt_kernel_installed_info gauge',
    ]
    for release, version in sorted(kernels.items()):
        kernel_info_lines.append(
            f'apt_kernel_installed_info{{release="{esc(release)}",version="{esc(version)}"}} 1'
        )
    kernel_info_block = '\n'.join(kernel_info_lines) + '\n'

    reboot = reboot_required(os.uname().release, kernels.keys())

    body = f'''# HELP apt_pending_updates Number of pending apt package updates.
# TYPE apt_pending_updates gauge
apt_pending_updates {total}
# HELP apt_pending_updates_patch Pending updates classified as patch (epoch/revision-only or third-component bump).
# TYPE apt_pending_updates_patch gauge
apt_pending_updates_patch {counts['patch']}
# HELP apt_pending_updates_minor Pending updates classified as minor (second-component bump).
# TYPE apt_pending_updates_minor gauge
apt_pending_updates_minor {counts['minor']}
# HELP apt_pending_updates_major Pending updates classified as major (first-component bump, or unparseable version).
# TYPE apt_pending_updates_major gauge
apt_pending_updates_major {counts['major']}
# HELP apt_pending_kernel_update 1 if a kernel package is among the pending updates, else 0.
# TYPE apt_pending_kernel_update gauge
apt_pending_kernel_update {kernel}
# HELP apt_kernel_reboot_required 1 if a newer kernel image is installed than the running one (reboot needed), else 0.
# TYPE apt_kernel_reboot_required gauge
apt_kernel_reboot_required {reboot}
# HELP apt_updates_last_check_timestamp Unix time of last apt-updates-collector run.
# TYPE apt_updates_last_check_timestamp gauge
apt_updates_last_check_timestamp {ts}
{kernel_info_block}{info_block}'''

    os.makedirs(OUT_DIR, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=OUT_DIR, prefix='apt_updates.', suffix='.prom.tmp')
    try:
        with os.fdopen(fd, 'w') as f:
            f.write(body)
        os.chmod(tmp, 0o644)
        os.replace(tmp, OUT_PATH)
    except Exception:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise


if __name__ == '__main__':
    main()
