cinnamon-launcher

#!/usr/bin/python3
#
""" Launches or restarts cinnamon
"""

import os
import sys
import time
import gettext
from setproctitle import setproctitle
from pathlib import Path
import psutil
import subprocess
import threading
import gi
import shutil
gi.require_version('Gtk', '3.0')  # noqa
from gi.repository import Gtk, GLib, Gio, GLib

FALLBACK_COMMAND = "metacity"
FALLBACK_ARGS = ("--replace",)

LAUNCHER_BUS_NAME = "org.cinnamon.Launcher"
LAUNCHER_BUS_PATH = "/org/cinnamon/Launcher"

INTERFACE_XML = (
    "<node>"
    "  <interface name='org.cinnamon.Launcher'>"
    "    <method name='CanRestart'>"
    "      <arg type='b' name='result' direction='out'/>"
    "    </method>"
    "    <method name='TryRestart'>"
    "      <arg type='b' name='result' direction='out'/>"
    "    </method>"
    "  </interface>"
    "</node>"
)

gettext.install("cinnamon", "/usr/share/locale")

panel_process_name = None
panel_cmd = None
if shutil.which("mate-panel"):
    panel_process_name = "mate-panel"
    panel_cmd = "mate-panel --replace &"
elif shutil.which("gnome-panel"):
    panel_process_name = "gnome-panel"
    panel_cmd = "gnome-panel --replace &"
elif shutil.which("tint2"):
    panel_process_name = "tint2"
    panel_cmd = "killall tint2; tint2 &"

# Used as a decorator to run things in the background
def async_function(func):
    def wrapper(*args, **kwargs):
        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
        thread.daemon = True
        thread.start()
        return thread
    return wrapper

# Used as a decorator to run things in the main loop, from another thread
def idle_function(func):
    def wrapper(*args):
        GLib.idle_add(func, *args)
    return wrapper

class Launcher:

    def __init__(self):
        self.fb_pid = 0
        self.kill_lock = threading.Lock()

        self.polkit_agent_proc = None
        self.nm_applet_proc = None
        self.automount_proc = None
        self.can_restart = False
        self.dialog = None

        self.cinnamon_pid = os.fork()
        if self.cinnamon_pid == 0:
            os.execvp("cinnamon", ("cinnamon", "--replace", ) + tuple(sys.argv[1:]))
        else:
            self.settings = Gio.Settings(schema="org.cinnamon.launcher")
            self.settings.connect("changed::memory-limit-enabled", self.on_settings_changed)
            self.killed = False
            if self.settings.get_boolean("memory-limit-enabled"):
                print("Cinnamon memory limit enabled: %d MB" % self.settings.get_int("memory-limit"))
                self.monitor_memory()
            self.setup_dbus()
            self.wait_for_process()
            Gtk.main()

    def on_settings_changed(self, settings, key):
        print("Memory profiler status changed, restarting Cinnamon.")
        self.restart_cinnamon()

    def setup_dbus(self):
        self.node_info = Gio.DBusNodeInfo.new_for_xml(INTERFACE_XML)
        Gio.bus_own_name(
            Gio.BusType.SESSION,
            LAUNCHER_BUS_NAME,
            Gio.BusNameOwnerFlags.NONE,
            self.on_bus_acquired,
            None,
            None
        )

    def on_bus_acquired(self, connection, name):
        connection.register_object(
            LAUNCHER_BUS_PATH,
            self.node_info.interfaces[0],
            self.on_method_call
        )

    def on_method_call(self, connection, sender, object_path,
                       interface_name, method_name, parameters, invocation):
        if method_name == "CanRestart":
            invocation.return_value(GLib.Variant("(b)", (self.can_restart,)))
        elif method_name == "TryRestart":
            result = False
            if self.can_restart and self.dialog is not None:
                GLib.idle_add(self.dialog.response, Gtk.ResponseType.YES)
                result = True
            invocation.return_value(GLib.Variant("(b)", (result,)))

    @async_function
    def wait_for_process(self):
        exit_status = os.waitpid(self.cinnamon_pid, 0)[1]
        if not os.WIFEXITED(exit_status):
            with self.kill_lock:
                if self.killed:
                    # Restart Cinnamon
                    self.restart_cinnamon()
                else:
                    self.fb_pid = os.fork()
                    if self.fb_pid == 0:
                        # Start the fallback panel
                        if panel_cmd is not None:
                            os.system(panel_cmd)
                        # Start the fallback window manager
                        os.execvp(FALLBACK_COMMAND, (FALLBACK_COMMAND,) + FALLBACK_ARGS)
                    else:
                        self.launch_fallback_helpers()
                        self.confirm_restart()
        else:
            Gtk.main_quit()

    def launch_fallback_helpers(self):
        agents = ("/usr/lib/policykit-1-gnome/polkit-gnome-authentication-agent-1",
                  "/usr/libexec/polkit-gnome-authentication-agent-1",
                  "/usr/libexec/polkit-mate-authentication-agent-1")

        for path in agents:
            if Path(path).exists():
                print(f"Launching for fallback session: {path}")
                self.polkit_agent_proc = subprocess.Popen([path])
                break

        # Start NM applet
        if shutil.which("nm-applet"):
            print(f"Launching for fallback session: nm-applet")
            self.nm_applet_proc = subprocess.Popen(["nm-applet"])

        if shutil.which("csd-automount"):
            print(f"Launching for fallback session: csd-automount")
            self.automount_proc = subprocess.Popen(["csd-automount"])

    def kill_fallback_helpers(self):
        if self.polkit_agent_proc is not None:
            print("Killing fallback polkit agent")
            self.polkit_agent_proc.terminate()
            try:
                self.polkit_agent_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.polkit_agent_proc.kill()
            self.polkit_agent_proc = None

        if self.nm_applet_proc is not None:
            print("Killing fallback nm-applet")
            self.nm_applet_proc.terminate()
            try:
                self.nm_applet_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.nm_applet_proc.kill()
            self.nm_applet_proc = None

        if self.automount_proc is not None:
            print("Killing fallback csd-automount")
            self.automount_proc.terminate()
            try:
                self.automount_proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.automount_proc.kill()
            self.automount_proc = None

    @async_function
    def monitor_memory(self):
        if psutil.pid_exists(self.cinnamon_pid):
            process = psutil.Process(self.cinnamon_pid)
            while True:
                delay = self.settings.get_int("check-frequency")
                limit = self.settings.get_int("memory-limit")
                time.sleep(delay)
                if not process.is_running():
                    break
                memory = ((process.memory_full_info().rss - process.memory_full_info().shared) / 1024 ** 2)
                if memory > limit:
                    print ("Memory threshold reached: %d" % memory)
                    self.log_kill(memory, limit, delay)
                    with self.kill_lock:
                        self.killed = True
                        process.kill()
                    break

    def restart_cinnamon(self):
        os.execvp(sys.argv[0], sys.argv)

    def log_kill(self, memory, limit, delay):
        try:
            directory = os.path.join(GLib.get_user_state_dir(), 'cinnamon')
        except AttributeError:
            directory = os.path.expanduser("~/.cinnamon")

        string = "%s - Usage: %d MB - Limit: %d MB - Freq: %s sec" % (time.strftime("%Y.%m.%d %H:%M:%S"), memory, limit, delay)
        if not os.path.exists(directory):
            os.mkdir(directory)
        with open(os.path.join(directory, "restarts.log"), "a") as log_file:
            print(string, file=log_file)

    @idle_function
    def confirm_restart(self):
        self.can_restart = True

        d = Gtk.MessageDialog(title=None, parent=None, flags=0, message_type=Gtk.MessageType.WARNING)
        d.add_buttons(_("No"),  Gtk.ResponseType.NO,
                      _("Yes"), Gtk.ResponseType.YES)
        d.set_keep_above(True)
        d.set_position(Gtk.WindowPosition.CENTER)
        d.set_title(_("Cinnamon just crashed"))
        box = d.get_content_area()
        checkbutton = Gtk.CheckButton(label=_("Disable downloaded applets, desklets and extensions"))
        d.set_markup("<span size='large'><b>%s</b></span>\n\n%s\n\n%s\n\n" %
                (
                    _("You are currently running in fallback mode."),
                    _("If you suspect the source of the crash is a local extension, applet or desklet, you can disable downloaded add-ons using the checkbox below."),
                    _("Do you want to restart Cinnamon?"),
                ))
        box.set_border_width(20)
        box.add(checkbutton)
        checkbutton.show_all()

        self.dialog = d
        resp = d.run()
        self.dialog = None
        self.can_restart = False
        d.destroy()
        if resp == Gtk.ResponseType.YES:
            if checkbutton.get_active():
                os.environ["CINNAMON_TROUBLESHOOT"] = "1"
            os.system("killall %s" % panel_process_name)
            self.kill_fallback_helpers()
            os.system("kill %d" % self.fb_pid)
            os.waitpid(self.fb_pid, 0)
            self.restart_cinnamon()
        else:
            Gtk.main_quit()

if __name__ == "__main__":
    setproctitle("cinnamon-launcher")
    Launcher()