charm.py

#!/usr/bin/env python3
# Copyright 2020 David Garcia
# See LICENSE file for licensing details.

from apt.progress.base import InstallProgress
import logging
import os
import shutil
import time

from jinja2 import Template
from ops.charm import CharmBase
from ops.framework import StoredState
from ops.main import main
from ops.model import (
    MaintenanceStatus,
    ActiveStatus,
    # BlockedStatus,
)
from utils import (
    service_stop,
    service_restart,
    install_apt,
    remove_apt,
    shell,
    upgrade_apt,
)


# from typing import Dict, Any
logger = logging.getLogger(__name__)

APT_PROXY_PATH = "/etc/apt/apt.conf.d/99-HIVE-apt-proxy"
APT_PROXY_TEMPLATE = "./templates/proxy"
APT_REQUIREMENTS = [
    "firefox",
    "indicator-applet-session",
    "libnotify-bin",
    "mate-desktop",
    "mate-applets",
    "mate-applet-brisk-menu",
    "mate-control-center",
    "mate-indicator-applet",
    "mate-notification-daemon",
    "mate-session-manager",
    "mate-terminal",
    "xrdp",
]
SNAP_INSTALLS = [
    #"code --classic",
]
POLKIT_TEMPLATE = "./templates/color.pkla"
POLKIT_PATH = "/etc/polkit-1/localauthority/50-local.d/color.pkla"
STARTWM_TEMPLATE = "./templates/startwm.sh"
STARTWM_PATH = "/etc/xrdp/startwm.sh"
# WM_COMMAND = "startxfce4" # xubuntu-desktop
# WM_COMMAND = "budgie-desktop" # budgie-desktop-environment
WM_COMMAND = "mate-session"  # mate-desktop


class VirtualPCCharm(CharmBase, InstallProgress):
    _stored = StoredState()

    def __init__(self, *args):
        super().__init__(*args)
        InstallProgress.__init__(self)

        self._stored.set_default()
        self.last_status_update = time.time()

        # Basic hooks
        self.framework.observe(self.on.install, self._on_install)
        self.framework.observe(self.on.start, self._on_start)
        self.framework.observe(self.on.stop, self._on_stop)
        self.framework.observe(self.on.config_changed, self._on_config_changed)
        self.framework.observe(self.on.update_status, self._on_update_status)

        # Actions hooks
        self.framework.observe(self.on["add-package"].action, self._add_package)
        self.framework.observe(self.on["add-snap"].action, self._add_snap)
        self.framework.observe(self.on["announce"].action, self._announce)
        self.framework.observe(self.on["reboot"].action, self._reboot)
        self.framework.observe(self.on["remove-package"].action, self._remove_package)
        self.framework.observe(self.on["remove-snap"].action, self._remove_snap)
        self.framework.observe(self.on["update-system"].action, self._update_system)

        # Relations hooks

    # Override InstallProgress to update our status
    def status_change(self, pkg, percent, status):
        if (time.time() - self.last_status_update) < 2:
            return
        self.last_status_update = time.time()
        message = str(int(percent)) + "% " + status
        self.unit.status = MaintenanceStatus(message)

    # Basic hooks
    def _on_install(self, _):

        self.unit.status = MaintenanceStatus("Setting up apt proxy")
        with open(APT_PROXY_TEMPLATE, "r") as template:
            content = Template(template.read()).render()
            with open(APT_PROXY_PATH, "w") as proxy:
                proxy.write(content)

        self.unit.status = MaintenanceStatus("Installing apt packages")
        install_apt(packages=APT_REQUIREMENTS, update=True, progress=self)
        service_stop('xrdp')

        self.unit.status = MaintenanceStatus("Installing snaps")
        for snap in SNAP_INSTALLS:
            shell("snap install " + snap)

        self.unit.status = MaintenanceStatus("Setting default display manager")
        shell("echo /usr/sbin/lightdm | sudo tee /etc/X11/default-display-manager")

        self.unit.status = MaintenanceStatus("Adding XRDP to ssl-cert group")
        shell("adduser xrdp ssl-cert")

        self.unit.status = MaintenanceStatus("Generating Window Manager startup script")
        with open(STARTWM_TEMPLATE, "r") as template:
            content = Template(template.read()).render(command=WM_COMMAND)
            with open(STARTWM_PATH, "w") as startwm:
                startwm.write(content)

        self.unit.status = MaintenanceStatus("Generating Polkit files")
        with open(POLKIT_TEMPLATE, "r") as template:
            content = Template(template.read()).render()
            with open(POLKIT_PATH, "w") as polkit:
                polkit.write(content)

        self._stored.installed = True
        self.unit.status = self._get_current_status()

    def _on_start(self, _):
        self.unit.status = MaintenanceStatus("Starting XRDP server")
        service_restart('xrdp')
        self._stored.started = True
        self.unit.status = self._get_current_status()

    def _on_stop(self, _):
        service_stop('xrdp')
        self._stored.started = False
        self.unit.status = self._get_current_status()

    def _on_config_changed(self, _):
        self.unit.status = self._get_current_status()

    def _on_update_status(self, _):
        self.unit.status = self._get_current_status()

    # Action hooks
    def _add_package(self, event):
        self.unit.status = MaintenanceStatus("Installing apt packages")
        install_apt(packages=event.params["package"].split(','),
                    update=True, progress=self)
        self.unit.status = self._get_current_status()

    def _add_snap(self, event):
        self.unit.status = MaintenanceStatus("Installing snaps")
        for snap in event.params["package"].split(','):
            shell("snap install " + snap)
        self.unit.status = self._get_current_status()

    def _announce(self, event):
        self.unit.status = MaintenanceStatus("Announce")
        message = event.params["message"]
        shell("su - ubuntu -c 'XDG_RUNTIME_DIR=/run/user/$(id -u) notify-send \"" + message + "\"'")
        self.unit.status = self._get_current_status()

    def _reboot(self, _):
        self.unit.status = MaintenanceStatus("Rebooting server")
        shell("su - ubuntu -c 'XDG_RUNTIME_DIR=/run/user/$(id -u) notify-send \"System is going down for reboot in 60 seconds\" -u critical'")
        shell("shutdown -r +1")
        self.unit.status = self._get_current_status()

    def _remove_package(self, event):
        self.unit.status = MaintenanceStatus("Removing apt packages")
        remove_apt(packages=event.params["package"].split(','),
                   update=True, progress=self)
        self.unit.status = self._get_current_status()

    def _remove_snap(self, event):
        self.unit.status = MaintenanceStatus("Removing snaps")
        for snap in event.params["package"].split(','):
            shell("snap remove " + snap)
        self.unit.status = self._get_current_status()

    def _update_system(self, _):
        self.unit.status = MaintenanceStatus("Updating system")
        upgrade_apt(update=True, progress=self)
        self.unit.status = self._get_current_status()

    # Relation hooks

    # Private functions
    def _get_current_status(self):
        status_type = ActiveStatus
        status_msg = ""
        if self._stored.installed:
            status_msg = "Ready"
        return status_type(status_msg)


if __name__ == "__main__":
    main(VirtualPCCharm)