#!/usr/bin/python3

import json
import os
import subprocess
import sys

ROLES="/etc/libvirt/privileges.d/roles"
ROLES_USER_DEFINED="/etc/libvirt/privileges.d/roles.user"
USERS="/etc/libvirt/privileges.d/users"
VMS="/etc/libvirt/privileges.d/vms"

LOG_HELPER="/usr/libexec/libvirt-log-helper"

class RoleModelGenericError(Exception):
    pass


class NoSuchRoleError(RoleModelGenericError):
    def __init__(self, role_name):
        self.error_msg = "No such role: %s" % role_name


class NoSuchUserError(RoleModelGenericError):
    def __init__(self, user_name):
        self.error_msg = "No such user: %s" % user_name


class NoSuchVmError(RoleModelGenericError):
    def __init__(self, vm_name):
        self.error_msg = "No such vm: %s" % vm_name


class DuplicateRoleError(RoleModelGenericError):
    def __init__(self, user_name, role_name):
        self.error_msg = "Role %s already assigned to user %s" % (user_name, role_name)


class NoSuchUserRoleError(RoleModelGenericError):
    def __init__(self, user_name, role_name):
        self.error_msg = "User %s is not associated with role %s" % (user_name, role_name)


class RoleModel:
    def __init__(self):
        with open(ROLES, "r") as roles_fp, open(USERS, "r") as users_fp, open(VMS, "r") as vms_fp:
            self.roles = json.load(roles_fp)
            self.users = json.load(users_fp)
            self.vms = json.load(vms_fp)
        self.roles_user = {}
        if os.path.exists(ROLES_USER_DEFINED):
            with open(ROLES_USER_DEFINED) as roles_user_fp:
                try:
                    self.roles_user = json.load(roles_user_fp)
                except json.decoder.JSONDecodeError:
                    print("File %s JSON Format error." % ROLES_USER_DEFINED)

    def get_role_list(self):
        all_roles = self.roles
        all_roles.update(self.roles_user)
        return [role for role in all_roles.keys()]

    def get_user_list(self):
        return [user for user in self.users.keys()]

    def get_user_roles(self, user):
        if user in self.users.keys():
            return self.users[user]
        else:
            raise NoSuchUserError(user)

    def get_vm_user_roles(self, vm):
        if vm in self.vms.keys():
            return self.vms[vm]
        else:
            raise NoSuchVmError(vm)

    def __check_role(self, role, user):
        return (role in self.get_role_list() or role in self.get_user_roles(user))

    def __check_user(self, user):
        return (user in self.get_user_list())

    def assign_role(self, user, role):
        if (not self.__check_role(role, user)):
            raise NoSuchRoleError(role)
        try:
            if (not role in self.users[user]):
                self.users[user].append(role)
            else:
                raise DuplicateRoleError(user, role)
        except KeyError:
            self.users[user] = [role]

    def revoke_role(self, user, role):
        if (not self.__check_role(role, user)):
            raise NoSuchRoleError(role)
        if (not self.__check_user(user)):
            raise NoSuchUserError(user)
        if (role in self.users[user]):
            self.users[user].remove(role)
        else:
            raise NoSuchUserRoleError(user, role)

    def assign_role_to_vm(self, user, role, vm):
        if (not self.__check_role(role, user)):
            raise NoSuchRoleError(role)
        try:
            vm_associated_roles = self.vms[vm]
            try:
                if (not role in self.vms[vm][user]):
                    self.vms[vm][user].append(role)
                else:
                    raise DuplicateRoleError(user, role)
            except KeyError:
                self.vms[vm][user] = list([role])
        except KeyError:
            self.vms[vm] = {user: list([role])}

    def revoke_role_from_vm(self, user, role, vm):
        if (not self.__check_role(role, user)):
            raise NoSuchRoleError(role)
        try:
            vm_associated_roles = self.vms[vm]
            try:
                if (role in vm_associated_roles[user]):
                    self.vms[vm][user].remove(role)
                else:
                    raise NoSuchUserRoleError(user, role)
            except KeyError:
                print("User %s have no associated roles with vm %s" % (user, vm))
        except KeyError:
            print("No associated user roles with vm %s" % vm)

    def save(self):
        with open(USERS, "w") as users_fp, open(VMS, "w") as vms_fp:
            json.dump(self.users, users_fp, indent=2, separators=(',', ': '))
            json.dump(self.vms, vms_fp, indent=2, separators=(',', ': '))


if __name__ == "__main__":
    if os.getuid() != 0:
        print("Only root allowed to use this tool!")
        exit(1)
    if len(sys.argv) == 4:
        action = sys.argv[1]
        user = sys.argv[2]
        role = sys.argv[3]
    elif len(sys.argv) == 5:
        action = sys.argv[1]
        user = sys.argv[2]
        role = sys.argv[3]
        vm = sys.argv[4]
    elif len(sys.argv) == 3:
        action = sys.argv[1]
        obj = sys.argv[2]
    else:
        print("Argc is invalid")
        exit(2)
    role_model = RoleModel()
    if action == "assign-role":
        try:
            role_model.assign_role(user, role)
            role_model.save()
            subprocess.run([LOG_HELPER, "RoleModel", "--user=%s" % user, "--role=%s" % role, "--action=%s" % action])
            print("Role %s assigned to user %s" % (role, user))
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    elif action == "revoke-role":
        try:
            role_model.revoke_role(user, role)
            role_model.save()
            subprocess.run([LOG_HELPER, "RoleModel", "--user=%s" % user, "--role=%s" % role, "--action=%s" % action])
            print("Role %s revoked from user %s" % (role, user))
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    elif action == "assign-role-to-vm":
        try:
            role_model.assign_role_to_vm(user, role, vm)
            role_model.save()
            subprocess.run([LOG_HELPER, "RoleModel", "--vm=%s" % vm, "--user=%s" % user,
                            "--role=%s" % role, "--action=%s" % action])
            print("Role %s assigned to user %s for vm (%s)" % (role, user, vm))
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    elif action == "revoke-role-from-vm":
        try:
            role_model.revoke_role_from_vm(user, role, vm)
            role_model.save()
            subprocess.run([LOG_HELPER, "RoleModel", "--vm=%s" % vm, "--user=%s" % user,
                            "--role=%s" % role, "--action=%s" % action])
            print("Role %s revoked from user %s for vm (%s)" % (role, user, vm))
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    elif action == "list-user-roles":
        try:
            user_roles = role_model.get_user_roles(obj)
            print("User %s associated with following roles:" % obj)
            for role in user_roles:
                print("\t%s" % role)
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    elif action == "list-vm-roles":
        try:
            vm_users_roles = role_model.get_vm_user_roles(obj)
            print("VM %s have the following associated users with roles:" % obj)
            for user_roles_key in vm_users_roles.keys():
                print("\t%s: " % user_roles_key)
                for user_role in vm_users_roles[user_roles_key]:
                    print("\t\t%s" % user_role)
        except (RoleModelGenericError) as e:
            print(e.error_msg)
    else:
        print("Unknown action")