#!/usr/bin/python3 -E


from __future__ import print_function
import os
import errno
import shutil
import sys
from optparse import OptionParser


try:
    import selinux
    import semanage
except ImportError:
    print("You must install libselinux-python and libsemanage-python before running this tool", file=sys.stderr)
    exit(1)


def copy_file(src, dst):
    if DEBUG:
        print("copying %s to %s" % (src, dst))
    try:
        shutil.copy(src, dst)
    except OSError as the_err:
        (err, strerr) = the_err.args
        print("Could not copy %s to %s, %s" % (src, dst, strerr), file=sys.stderr)
        exit(1)


def create_dir(dst, mode):
    if DEBUG:
        print("Making directory %s" % dst)
    try:
        os.makedirs(dst, mode)
    except OSError as the_err:
        (err, stderr) = the_err.args
        if err == errno.EEXIST:
            pass
        else:
            print("Error creating %s" % dst, file=sys.stderr)
            exit(1)


def create_file(dst):
    if DEBUG:
        print("Making file %s" % dst)
    try:
        open(dst, 'a').close()
    except OSError as the_err:
        (err, stderr) = the_err.args
        print("Error creating %s" % dst, file=sys.stderr)
        exit(1)


def copy_module(store, name, base):
    if DEBUG:
        print("Install module %s" % name)
    (file, ext) = os.path.splitext(name)
    if ext != ".pp":
        # Stray non-pp file in modules directory, skip
        print("warning: %s has invalid extension, skipping" % name, file=sys.stderr)
        return
    try:
        if base:
            root = oldstore_path(store)
        else:
            root = oldmodules_path(store)

        bottomdir = bottomdir_path(store)

        os.mkdir("%s/%s" % (bottomdir, file))

        copy_file(os.path.join(root, name), "%s/%s/hll" % (bottomdir, file))

        # This is the ext file that will eventually be used to choose a compiler
        efile = open("%s/%s/lang_ext" % (bottomdir, file), "w+", 0o600)
        efile.write("pp")
        efile.close()

    except (IOError, OSError):
        print("Error installing module %s" % name, file=sys.stderr)
        exit(1)


def disable_module(file, name, disabledmodules):
    if DEBUG:
        print("Disabling %s" % name)
    (disabledname, disabledext) = os.path.splitext(file)
    create_file("%s/%s" % (disabledmodules, disabledname))


def migrate_store(store):
    oldstore = oldstore_path(store)
    oldmodules = oldmodules_path(store)
    disabledmodules = disabledmodules_path(store)
    newstore = newstore_path(store)
    newmodules = newmodules_path(store)
    bottomdir = bottomdir_path(store)

    print("Migrating from %s to %s" % (oldstore, newstore))

    # Build up new directory structure
    create_dir("%s/%s" % (newroot_path(), store), 0o755)
    create_dir(newstore, 0o700)
    create_dir(newmodules, 0o700)
    create_dir(bottomdir, 0o700)
    create_dir(disabledmodules, 0o700)

    # Special case for base since it was in a different location
    copy_module(store, "base.pp", 1)

    # Dir structure built, start copying files
    for root, dirs, files in os.walk(oldstore):
        if root == oldstore:
            # This is the top level directory, need to move
            for name in files:
                # Check to see if it is in TOPPATHS and copy if so
                if name in TOPPATHS:
                    if name == "seusers":
                        newname = "seusers.local"
                    else:
                        newname = name
                    copy_file(os.path.join(root, name), os.path.join(newstore, newname))

        elif root == oldmodules:
            # This should be the modules directory
            for name in files:
                (file, ext) = os.path.splitext(name)
                if name == "base.pp":
                    print("Error installing module %s, name conflicts with base" % name, file=sys.stderr)
                    exit(1)
                elif ext == ".disabled":
                    disable_module(file, name, disabledmodules)
                else:
                    copy_module(store, name, 0)


def rebuild_policy():
    # Ok, the modules are loaded, lets try to rebuild the policy
    print("Attempting to rebuild policy from %s" % newroot_path())

    curstore = selinux.selinux_getpolicytype()[1]

    handle = semanage.semanage_handle_create()
    if not handle:
        print("Could not create semanage handle", file=sys.stderr)
        exit(1)

    semanage.semanage_select_store(handle, curstore, semanage.SEMANAGE_CON_DIRECT)

    if not semanage.semanage_is_managed(handle):
        semanage.semanage_handle_destroy(handle)
        print("SELinux policy is not managed or store cannot be accessed.", file=sys.stderr)
        exit(1)

    rc = semanage.semanage_access_check(handle)
    if rc < semanage.SEMANAGE_CAN_WRITE:
        semanage.semanage_handle_destroy(handle)
        print("Cannot write to policy store.", file=sys.stderr)
        exit(1)

    rc = semanage.semanage_connect(handle)
    if rc < 0:
        semanage.semanage_handle_destroy(handle)
        print("Could not establish semanage connection", file=sys.stderr)
        exit(1)

    semanage.semanage_set_rebuild(handle, 1)

    rc = semanage.semanage_begin_transaction(handle)
    if rc < 0:
        semanage.semanage_handle_destroy(handle)
        print("Could not begin transaction", file=sys.stderr)
        exit(1)

    rc = semanage.semanage_commit(handle)
    if rc < 0:
        print("Could not commit transaction", file=sys.stderr)

    semanage.semanage_handle_destroy(handle)


def oldroot_path():
    return "%s/etc/selinux" % ROOT


def oldstore_path(store):
    return "%s/%s/modules/active" % (oldroot_path(), store)


def oldmodules_path(store):
    return "%s/modules" % oldstore_path(store)


def disabledmodules_path(store):
    return "%s/disabled" % newmodules_path(store)


def newroot_path():
    return "%s%s" % (ROOT, PATH)


def newstore_path(store):
    return "%s/%s/active" % (newroot_path(), store)


def newmodules_path(store):
    return "%s/modules" % newstore_path(store)


def bottomdir_path(store):
    return "%s/%s" % (newmodules_path(store), PRIORITY)


if __name__ == "__main__":

    parser = OptionParser()
    parser.add_option("-p", "--priority", dest="priority", default="100",
                      help="Set priority of modules in new store (default: 100)")
    parser.add_option("-s", "--store", dest="store", default=None,
                      help="Store to read from and write to")
    parser.add_option("-d", "--debug", dest="debug", action="store_true", default=False,
                      help="Output debug information")
    parser.add_option("-c", "--clean", dest="clean", action="store_true", default=False,
                      help="Clean old modules directory after migrate (default: no)")
    parser.add_option("-n", "--norebuild", dest="norebuild", action="store_true", default=False,
                      help="Disable rebuilding policy after migration (default: no)")
    parser.add_option("-P", "--path", dest="path",
                      help="Set path for the policy store (default: /var/lib/selinux)")
    parser.add_option("-r", "--root", dest="root",
                      help="Set an alternative root for the migration (default: /)")

    (options, args) = parser.parse_args()

    DEBUG = options.debug
    PRIORITY = options.priority
    TYPE = options.store
    CLEAN = options.clean
    NOREBUILD = options.norebuild
    PATH = options.path
    if PATH is None:
        PATH = "/var/lib/selinux"

    ROOT = options.root
    if ROOT is None:
        ROOT = ""

    # List of paths that go in the active 'root'
    TOPPATHS = [
        "commit_num",
        "ports.local",
        "interfaces.local",
        "nodes.local",
        "booleans.local",
        "file_contexts.local",
        "seusers",
        "users.local",
        "users_extra",
        "users_extra.local",
        "disable_dontaudit",
        "preserve_tunables",
        "policy.kern",
        "file_contexts",
        "homedir_template",
        "pkeys.local",
        "ibendports.local"]

    create_dir(newroot_path(), 0o755)

    stores = None
    if TYPE is not None:
        stores = [TYPE]
    else:
        stores = os.listdir(oldroot_path())

    # find stores in oldroot and migrate them to newroot if necessary
    for store in stores:
        if not os.path.isdir(oldmodules_path(store)):
            # already migrated or not an selinux store
            continue

        if os.path.isdir(newstore_path(store)):
            # store has already been migrated, but old modules dir still exits
            print("warning: Policy type %s has already been migrated, but modules still exist in the old store. Skipping store." % store, file=sys.stderr)
            continue

        migrate_store(store)

        if CLEAN is True:
            def remove_error(function, path, execinfo):
                print("warning: Unable to remove old store modules directory %s. Cleaning failed." % oldmodules_path(store), file=sys.stderr)
            shutil.rmtree(oldmodules_path(store), onerror=remove_error)

    if NOREBUILD is False:
        rebuild_policy()