Spotter-VM/basic/srv/vm/mgr/pkgmgr.py

# -*- coding: utf-8 -*-

import hashlib
import json
import os
import requests
import shutil
import subprocess
import tempfile

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.serialization import load_pem_public_key

PUB_FILE = '/srv/vm/packages.pub'
LXC_ROOT = '/var/lib/lxc'

class PackageManager:
    def __init__(self, conf):
        # Load JSON configuration
        self.conf = conf
        self.online_packages = {}
        self.pending = 0

    def get_repo_resource(self, url, stream=False):
        return requests.get('{}/{}'.format(self.conf['repo']['url'], url), auth=(self.conf['repo']['user'], self.conf['repo']['pwd']), stream=stream)

    def fetch_online_packages(self):
        # Fetches and verifies online packages. Can raise InvalidSignature
        packages = self.get_repo_resource('packages').content
        packages_sig = self.get_repo_resource('packages.sig').content
        with open(PUB_FILE, 'rb') as f:
            pub_key = load_pem_public_key(f.read(), default_backend())
        pub_key.verify(packages_sig, packages, ec.ECDSA(hashes.SHA512()))
        self.online_packages = json.loads(packages)

    def register_pending_installation(self, name):
        # Registers pending installation. Fetch online packages here instead of install_pacakges() to fail early if the repo isn't reachable
        self.fetch_online_packages()
        self.pending = 1
        # Return total size for download
        deps = [d for d in self.get_install_deps(name) if d not in self.conf['packages']]
        return sum(self.online_packages[d]['size'] for d in deps)

    def install_package(self, name):
        # Main installation function. Wrapper for download, registration and install script
        deps = [d for d in self.get_install_deps(name) if d not in self.conf['packages']]
        for dep in deps:
            self.download_package(dep)
            self.register_package(dep)
            self.run_install_script(dep)
        self.pending = 0

    def uninstall_package(self, name):
        # Main uninstallation function. Wrapper for uninstall script, filesystem purge and unregistration
        deps = self.get_install_deps(name, False)[::-1]
        for dep in deps:
            if dep not in self.get_uninstall_deps():
                self.run_uninstall_script(dep)
                self.purge_package(dep)
                self.unregister_package(dep)

    def download_package(self, name):
        # Downloads, verifies, unpacks and sets up a package
        tmp_archive = tempfile.mkstemp('.tar.xz')[1]
        r = self.get_repo_resource('{}.tar.xz'.format(name), True)
        with open(tmp_archive, 'wb') as f:
            for chunk in r.iter_content(chunk_size=65536):
                if chunk:
                    self.pending += f.write(chunk)
        # Verify hash
        if self.online_packages[name]['sha512'] != hash_file(tmp_archive):
            raise InvalidSignature(name)
        # Unpack
        subprocess.run(['tar', 'xJf', tmp_archive], cwd='/')
        os.unlink(tmp_archive)

    def purge_package(self, name):
        # Removes package and shared data from filesystem
        shutil.rmtree(os.path.join(LXC_ROOT, self.conf['packages'][name]['lxcpath']))
        srv_dir = os.path.join('/srv/', name)
        if os.path.exists(srv_dir):
            shutil.rmtree(srv_dir)

    def run_install_script(self, name):
        # Runs install.sh for a package, if the script is present
        install_dir = os.path.join('/srv/', name, 'install')
        install_script = os.path.join('/srv/', name, 'install.sh')
        if os.path.exists(install_script):
            subprocess.run(install_script)
            os.unlink(install_script)
        if os.path.exists(install_dir):
            shutil.rmtree(install_dir)

    def run_uninstall_script(self, name):
        # Runs uninstall.sh for a package, if the script is present
        uninstall_script = os.path.join('/srv/', name, 'uninstall.sh')
        if os.path.exists(uninstall_script):
            subprocess.run(uninstall_script)

    def register_package(self, name):
        # Registers a package in local configuration
        metadata = self.online_packages[name]
        self.conf['packages'][name] = {
            'deps': metadata['deps'],
            'lxcpath': metadata['lxcpath'],
            'version': metadata['version']
        }
        # If host definition is present, register the package as application
        if 'host' in metadata:
            self.conf['apps'][name] = {
                'title': metadata['title'],
                'host': metadata['host'],
                'login': 'N/A',
                'password': 'N/A',
                'visible': False
            }
        self.conf.save()

    def unregister_package(self, name):
        # Removes a package from local configuration
        del self.conf['packages'][name]
        if name in self.conf['apps']:
            del self.conf['apps'][name]
        self.conf.save()

    def get_install_deps(self, name, online=True):
        # Flatten dependency tree for a package while preserving the dependency order
        packages = self.online_packages if online else self.conf['packages']
        deps = packages[name]['deps'].copy()
        for dep in deps[::-1]:
            deps[:0] = [d for d in self.get_install_deps(dep, online)]
        deps = list(dict.fromkeys(deps + [name]))
        return deps

    def get_uninstall_deps(self):
        # Create reverse dependency tree for all installed packages
        deps = {}
        for pkg in self.conf['packages']:
            for d in self.conf['packages'][pkg]['deps']:
                deps.setdefault(d, []).append(pkg)
        return deps

def hash_file(file_path):
    sha512 = hashlib.sha512()
    with open(file_path, 'rb') as f:
        while True:
            data = f.read(65536)
            if not data:
                break
            sha512.update(data)
    return sha512.hexdigest()