steam-backlog-enforcer/steam_backlog_enforcer/_playtime_procs.py

"""Identifying gaming processes by their /proc cmdline.

Process-name matching alone misses launchers run through an interpreter
(``lutris`` is a Python script; several Minecraft launchers are ``java -jar``),
so this reads argv and looks past the interpreter to the real program.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import TYPE_CHECKING

from steam_backlog_enforcer._counted_procs import key_by_name, load_counted_processes
from steam_backlog_enforcer._total_block_launchers import LAUNCHER_PROCESS_NAMES
from steam_backlog_enforcer.enforcer import (
    get_pids_by_process_names,
    get_running_steam_game_pids,
)

if TYPE_CHECKING:
    from steam_backlog_enforcer._playtime_state import PlaytimeRules

_PROC = Path("/proc")

# argv[0] basenames that tell us nothing: the real program is argv[1].
_INTERPRETERS = frozenset({"bash", "env", "java", "perl", "python", "python3", "sh"})

# An interpreter invocation needs at least `<interp> <script>` to name a program.
_MIN_ARGV_FOR_INTERPRETER = 2


def get_pids_by_cmdline_names(names: frozenset[str]) -> dict[int, str]:
    """Scan ``/proc/*/cmdline`` for processes whose program name is in *names*.

    Complements ``enforcer.get_pids_by_process_names``, which matches on
    ``comm`` and therefore cannot see interpreter-launched programs: the kernel
    records ``/usr/bin/lutris`` as ``python3``. When ``argv[0]``'s basename is a
    known interpreter this falls through to ``argv[1]``.

    The daemon itself runs under ``python3``, so its own PID is skipped.

    Args:
        names: Program basenames to match.

    Returns:
        Mapping of PID to the matched name.
    """
    own_pid = os.getpid()
    found: dict[int, str] = {}

    for entry in _PROC.iterdir():
        if not entry.name.isdigit():
            continue
        pid = int(entry.name)
        if pid == own_pid:
            continue
        matched = _match_cmdline(entry, names)
        if matched is not None:
            found[pid] = matched

    return found


def _match_cmdline(entry: Path, names: frozenset[str]) -> str | None:
    """Return the name in *names* that *entry*'s cmdline runs, if any.

    Args:
        entry: A ``/proc/<pid>`` directory.
        names: Program basenames to match.

    Returns:
        The matched name, or ``None``.
    """
    try:
        raw = (entry / "cmdline").read_bytes()
    except OSError, ValueError:
        return None

    argv = [
        part for part in raw.decode("utf-8", errors="replace").split("\x00") if part
    ]
    if not argv:
        return None

    first = Path(argv[0]).name
    if first in names:
        return first
    if first not in _INTERPRETERS or len(argv) < _MIN_ARGV_FOR_INTERPRETER:
        return None
    second = Path(argv[1]).name
    return second if second in names else None


def process_name(pid: int) -> str | None:
    """Return *pid*'s program name, or ``None`` if it is no longer running.

    Existence is the point as much as the name: callers resolve PIDs recorded
    earlier, and a PID that has since been recycled would otherwise be reported
    under whatever process now holds the number.

    Args:
        pid: Process id to look up.

    Returns:
        The contents of ``/proc/<pid>/comm``, or ``None``.
    """
    try:
        return (_PROC / str(pid) / "comm").read_text(encoding="utf-8").strip() or None
    except OSError, ValueError:
        return None


def _merge_named(found: dict[int, str], keys: dict[str, str]) -> None:
    """Add name-matched PIDs to *found* without displacing existing keys.

    ``setdefault`` is the point: a Steam game already attributed by its
    ``SteamAppId`` must keep that key even if its ``comm`` also happens to match
    a launcher name.

    A scanner can only return names it was asked for, so an unmapped name is
    unreachable in production — it is skipped rather than raising, because the
    alternative is a ``KeyError`` inside the enforce loop.

    Args:
        found: Accumulating PID-to-key mapping, mutated in place.
        keys: Program basename to attribution key.
    """
    names = frozenset(keys)
    for source in (get_pids_by_process_names(names), get_pids_by_cmdline_names(names)):
        for pid, name in source.items():
            key = keys.get(name)
            if key is not None:
                found.setdefault(pid, key)


def qualifying_pids(rules: PlaytimeRules) -> dict[int, str]:
    """Return PIDs whose runtime counts against the daily budget, with owners.

    Steam games are identified by the ``SteamAppId`` environment variable, using
    the same ``!= 0`` predicate ``enforcer.enforce_allowed_game`` uses to exclude
    the Steam client tree — browsing the store is not gaming.

    The value is an attribution key (``app:<id>``, ``launcher:<name>`` or
    ``proc:<id>``) so that the budget can record *which* game it billed. The
    app id was previously destructured and dropped on the first line here, one
    step before accumulation.

    ``counted_processes`` is deliberately *not* gated on ``count_launchers``:
    that switch covers the sixteen hardcoded launchers, and turning it off must
    not silently stop billing a game the user explicitly listed.

    Args:
        rules: Policy for this tick.

    Returns:
        Mapping of qualifying PID to its attribution key.
    """
    found: dict[int, str] = {
        pid: f"app:{app_id}"
        for pid, app_id in get_running_steam_game_pids().items()
        if app_id != 0
    }
    if rules.count_launchers:
        _merge_named(
            found, {name: f"launcher:{name}" for name in LAUNCHER_PROCESS_NAMES}
        )

    counted = key_by_name(load_counted_processes())
    if counted:
        _merge_named(found, counted)
    return found


def attributed_key(qualifying: dict[int, str]) -> str:
    """Return the single attribution key this tick should be credited to.

    With exactly one game qualifying, it gets the credit. With two or more
    running at once there is no honest way to say which one earned the tick,
    so this returns ``""``: the tick still bills, and the gap shows up as
    Unattributed rather than being charged to a guess.

    Args:
        qualifying: Mapping of qualifying PID to attribution key.

    Returns:
        The attribution key, or ``""`` if none can be determined.
    """
    distinct = set(qualifying.values())
    return distinct.pop() if len(distinct) == 1 else ""