#!/bin/bash
# vim: et sts=4 sw=4

# SPDX-License-Identifier: MIT
# Copyright 2021 Offensive Security
# Copyright (c) Microsoft Corporation

# Credits:
# * https://github.com/mimura1133/linux-vm-tools
# * https://github.com/microsoft/linux-vm-tools
# * https://c-nergy.be/blog/?p=12429
# * https://c-nergy.be/blog/?p=11755
# * https://c-nergy.be/blog/?p=11336

set -e
set -u

BCK="kali-tweaks-orig"    # extension for backup files

exit_with_usage() {
    echo >&2 "Usage: $(basename $0) enable|disable|check"
    exit 1
}

assert_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo >&2 "This script must be run with root privileges"
        exit 1
    fi
}

assert_pkg_installed() {
    if ! dpkg -s $1 2>/dev/null | grep -q "ok installed"; then
	echo >&2 "Package '$1' must be installed"
	exit 1
    fi
}

backup_file() {
    [ -f "$1" ] || return 0
    cp -f --preserve=all "$1" "$1.$BCK"
}

restore_file() {
    [ -f "$1.$BCK" ] || return 0
    mv -f "$1.$BCK" "$1"
}

do_configure() {

    # Configure XRDP, cf. xrdp.ini(5).
    # * use vsock transport
    # * use rdp security
    # * remove encryption validation
    # * disable bitmap compression, since its local its much faster
    #
    # Note: there are several 'port=' statements in the config file,
    # we match 'port=3389' (the default value) to make sure to change
    # only this line.
    backup_file /etc/xrdp/xrdp.ini
    sed -i \
        -e 's|^ *port=3389|port=vsock://-1:3389|' \
        -e 's|^ *security_layer=.*|security_layer=rdp|' \
        -e 's|^ *crypt_level=.*|crypt_level=none|' \
        -e 's|^ *bitmap_compression=.*|bitmap_compression=false|' \
        /etc/xrdp/xrdp.ini

    # Configure the XRDP session manager, cf. sesman.ini(5).
    # * set the first X display number available for xrdp-sesman to 0
    # * rename the redirected drives to 'shared-drives'.
    backup_file /etc/xrdp/sesman.ini
    sed -i \
        -e 's|^ *X11DisplayOffset=.*|X11DisplayOffset=0|' \
        -e 's|^ *FuseMountName=.*|FuseMountName=shared-drives|' \
        /etc/xrdp/sesman.ini

    # Add polkit-1 rules for colord access and repo refresh
    # cf. https://src.fedoraproject.org/rpms/xrdp/blob/rawhide/f/xrdp-polkit-1.rules
    cat > /etc/polkit-1/rules.d/xrdp-polkit-1.rules << 'EOF'
polkit.addRule(function(action, subject) {
    if ((action.id == "org.freedesktop.color-manager.create-device" ||
         action.id == "org.freedesktop.color-manager.create-profile"||
         action.id == "org.freedesktop.color-manager.delete-device" ||
         action.id == "org.freedesktop.color-manager.delete-profile" ||
         action.id == "org.freedesktop.color-manager.modify-device" ||
         action.id == "org.freedesktop.color-manager.modify-profile" ||
         action.id == "org.freedesktop.NetworkManager.network-control" ||
         action.id == "org.freedesktop.packagekit.system-sources-refresh") &&
        subject.active == true) {
            return polkit.Result.YES;
    }

    return polkit.Result.NOT_HANDLED;
});
EOF

    # Ensure the Hyper-V sockets module gets loaded.
    echo "hv_sock" > /etc/modules-load.d/hv_sock.conf
    if [ -d /run/systemd/system ]; then
        systemctl restart systemd-modules-load.service
    fi

    # GDM3 support (experimental)
    # GNOME users: you might want to uncomment that:
    #if [ -e /etc/gdm3/daemon.conf ]; then
    #    sed -i 's/#WaylandEnable=.*/WaylandEnable=false/' /etc/gdm3/daemon.conf
    #fi

}

do_unconfigure() {

    # Restore xrdp config files
    restore_file /etc/xrdp/xrdp.ini
    restore_file /etc/xrdp/sesman.ini
   
    # Drop the polkit-1 rules
    rm -f /etc/polkit-1/rules.d/xrdp-polkit-1.rules

    # Drop the old polkit pkla (replaced by .rules above in kali-tweaks 2022.4)
    rm -f /etc/polkit-1/localauthority/50-local.d/45-allow-colord.pkla

    # Unload the Hyper-V sockets module, don't auto-load it
    rm -f /etc/modules-load.d/hv_sock.conf
    rmmod hv_sock || :

}

is_configured() {
    # If the module is loaded, we assume that the system is configured.
    cat /proc/modules | grep -q '^hv_sock '
}


## main


[ $# -eq 1 ] || exit_with_usage

assert_pkg_installed xrdp

case "$1" in
    (enable) assert_root ; do_configure ;;
    (disable) assert_root ; do_unconfigure ;;
    (check) is_configured ;;
    (*) exit_with_usage ;;
esac
