#!/usr/bin/python3
# -*- coding: UTF-8 -*-
"""
Perform reboot after upgrades as directed by bonding-setup.

This should normally only be called via a systemd service.
"""
# © 2020, Multapplied Networks, Inc.

import errno
import fcntl
import os
import subprocess
import sys
import time


class PlatformError(Exception):
    pass


class Platform:
    UNDETECTED_PACKAGE_MANAGER_DELAY = 60

    def __init__(self):
        self.os_release = self.get_os_release()

    def get_os_release(self):
        "Return a dict of values from /etc/os-release"
        d = {}
        with open("/etc/os-release") as f:
            for line in f.readlines():
                if "=" in line:
                    key, val = line.split("=", 1)
                    d[key] = val.strip('"\n')
        return d

    @property
    def package_manager(self):
        osid = self.os_release.get("ID")

        if osid == "debian":
            return "apt"
        elif osid in ("rhel", "centos"):
            return "dnf"
        elif "suse" in osid:
            return "zypper"

        return None

    @property
    def systemctl_path(self):
        "Get the systemctl path"
        for path in ("/bin/systemctl", "/usr/bin/systemctl"):
            if os.path.exists(path):
                return path
        raise PlatformError("Could not determine systemctl path.")

    def wait_for_lock(self, path):
        "Wait for lock using fcntl locking"
        try:
            fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW)
        except EnvironmentError as e:
            print("ERROR: Could not open %s: %s" % (path, e), file=sys.stderr)
            # If it's not there it's probably not running
            return

        try:
            fcntl.lockf(fd, fcntl.LOCK_EX)
        finally:
            os.close(fd)

    def is_apt_running(self):
        for entry in os.listdir("/proc"):
            exe_link = "/proc/%s/exe" % entry
            if entry.isdigit() and os.path.islink(exe_link):
                try:
                    if os.readlink(exe_link).split()[0] in ("/usr/bin/apt", "/usr/bin/apt-get"):
                        return True
                except EnvironmentError as e:
                    if e.errno == errno.ENOENT:
                        # Process no longer exists
                        continue
        return False

    def wait_for_apt(self):
        while self.is_apt_running():
            time.sleep(1)

    def wait_for_dnf(self):
        # dnf does not have it's own locking, but it performs operations in a
        # single RPM transaction, so wait on the RPM DB lock
        self.wait_for_lock("/var/lib/rpm/.dbenv.lock")

    def is_zypper_running(self):
        try:
            with open("/var/run/zypp.pid", "r") as f:
                pid_data = f.read().strip()
        except EnvironmentError as e:
            if e.errno == errno.ENOENT:
                # Not present
                return False
            raise

        if not pid_data:
            # Not running
            return False

        pid = int(pid_data)

        # Check to see if PID is really zypper
        try:
            if os.readlink("/proc/%s/exe" % pid).endswith("/zypper"):
                return True
            else:
                # PID is some other program
                return False
        except EnvironmentError as e:
            if e.errno == errno.ENOENT:
                # Process no longer exists
                return False
            raise

    def wait_for_zypper(self):
        # zypper makes a series of RPM transactions and uses a PID file for
        # locking. However, it seems the locking is only used to read/write
        # the file, and not held during the process. Inspecting the process
        # for the PID seems to be what it actually does to avoid concurrent
        # execution
        while self.is_zypper_running():
            time.sleep(1)

    def wait_for_package_manager(self):
        """Wait for the package manager to be finished. If it is not detected,
        simply wait for a pre-defined period of time.
        """
        if self.package_manager == "apt":
            self.wait_for_apt()
        elif self.package_manager == "dnf":
            self.wait_for_dnf()
        elif self.package_manager == "zypper":
            self.wait_for_zypper()
        else:
            print("WARNING: Could not determine package manager. Waiting for %s seconds to allow for completion." %  self.UNDETECTED_PACKAGE_MANAGER_DELAY)
            time.sleep(self.UNDETECTED_PACKAGE_MANAGER_DELAY)

    def reboot(self):
        # Reboot via command line since we are likely in an upgrade scenario
        # and may not be able to rely on DBUS
        try:
            subprocess.check_call([self.systemctl_path, "reboot"])
        except subprocess.CalledProcessError as e:
            raise PlatformError("Could not execute systemctl to reboot: %s" % e.output)

if __name__ == "__main__":
    platform = Platform()
    platform.wait_for_package_manager()
    platform.reboot()
