nohang/nohang_notify_helper

#!/usr/bin/env python3

# print('Starting nohang_notify_helper')


def write(path, string):
    """
    """
    with open(path, 'w') as f:
        f.write(string)


def rline1(path):
    """read 1st line from path."""
    try:
        with open(path) as f:
            for line in f:
                return line
    except OSError:
        exit(1)


def rfile(path):
    """read file."""
    with open(path) as f:
        return f.read()


def re_pid_environ(pid):
    """
    read environ of 1 process
    returns tuple with USER, DBUS, DISPLAY like follow:
    ('user', 'DISPLAY=:0',
     'DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/1000/bus')
    returns None if these vars is not in /proc/[pid]/environ
    """
    try:
        env = str(rline1('/proc/' + pid + '/environ'))
        if display_env in env and dbus_env in env and user_env in env:
            env_list = env.split('\x00')

            # iterating over a list of process environment variables
            for i in env_list:
                if i.startswith(user_env):
                    user = i
                    if user == 'root':
                        return None
                    continue

                if i.startswith(display_env):
                    display = i[:10]
                    continue

                if i.startswith(dbus_env):
                    dbus = i
                    continue

                if i.startswith('HOME='):
                    # exclude Display Manager's user
                    if i.startswith('HOME=/var'):
                        return None

            try:
                env = user.partition('USER=')[2], display, dbus
            except UnboundLocalError:
                # print('notify helper: UnboundLocalError')
                return None

            return env

    except FileNotFoundError:
        # print('notify helper: FileNotFoundError')
        return None
    except ProcessLookupError:
        # print('notify helper: ProcessLookupError')
        return None


def root_notify_env():
    """return set(user, display, dbus)"""
    unsorted_envs_list = []
    # iterates over processes, find processes with suitable env
    for pid in listdir('/proc'):

        if path.exists('/proc/' + pid + '/exe') is True:
            one_env = re_pid_environ(pid)
            unsorted_envs_list.append(one_env)

    env = set(unsorted_envs_list)
    env.discard(None)

    # deduplicate dbus
    new_env = []
    end = []
    for i in env:
        key = i[0] + i[1]
        if key not in end:
            end.append(key)
            new_env.append(i)
        else:
            continue

    return new_env


try:
    write('/proc/self/oom_score_adj', '0')
except Exception:
    pass


try:
    from os import listdir, path, remove
    from subprocess import Popen, TimeoutExpired
    from sys import argv
except OSError:
    exit(1)


if len(argv) != 5:
    print('nohang_notify_helper: invalid input')
    exit()


with open('/proc/meminfo') as f:
    for line in f:
        if line.startswith('SwapTotal'):
            swap_total = int(line.split(':')[1][:-4])
            if swap_total > 0:
                wait_time = 10
            else:
                wait_time = 2

print('nohang_notify_helper: wait_time:', wait_time)

split_by = '#' * 16


uid = argv[2]

timestamp = argv[4]


path_to_cache = '/dev/shm/nohang_notify_cache_uid{}_time{}'.format(
    uid, timestamp
)


try:
    title, body = rfile(path_to_cache).split(split_by)
except FileNotFoundError:
    print('nohang_notify_helper: FileNotFoundError')
    exit(1)

remove(path_to_cache)

if uid != '0':
    cmd = ['notify-send', '--icon=dialog-warning', title, body]
    print('nohang_notify_helper: run cmd:', cmd)
    with Popen(cmd) as proc:
        try:
            proc.wait(timeout=wait_time)
        except TimeoutExpired:
            proc.kill()
            print('nohang_notify_helper: TimeoutExpired')
    exit()

display_env = 'DISPLAY='
dbus_env = 'DBUS_SESSION_BUS_ADDRESS='
user_env = 'USER='

list_with_envs = root_notify_env()
list_len = len(list_with_envs)

# if somebody logged in with GUI
if list_len > 0:

    for i in list_with_envs:
        print('Send GUI notification:', title, body, i)

    # iterating over logged-in users
    for i in list_with_envs:
        username, display_env, dbus_env = i[0], i[1], i[2]
        display_tuple = display_env.partition('=')
        dbus_tuple = dbus_env.partition('=')
        display_value = display_tuple[2]
        dbus_value = dbus_tuple[2]

        try:
            with Popen([
                'sudo', '-u', username,
                'env',
                'DISPLAY=' + display_value,
                'DBUS_SESSION_BUS_ADDRESS=' + dbus_value,
                'notify-send',
                '--icon=dialog-warning',
                title,
                body
            ]) as proc:
                try:
                    proc.wait(timeout=wait_time)
                except TimeoutExpired:
                    proc.kill()
                    print('TimeoutExpired: notify user: ' + username)
        except BlockingIOError:
            print('nohang_notify_helper: BlockingIOError')
        except OSError:
            print('nohang_notify_helper: OSError')
        except Exception:
            print('nohang_notify_helper: CANNOT SPAWN NOTIFY-SEND PROCESS')
else:
    print(
        'Not send GUI notification: [',
        title,
        body,
        ']. Nobody logged-in with GUI. Nothing to do.')