isobox

#!/usr/bin/python3

import sys
import os
import argparse
import traceback
import json
from libisobox.mnt import *
from libisobox.rootfs import *
from libisobox.unsquash_filesystem import *
from libisobox.chroot import *
from libisobox.filesystems import *
from libisobox.isobox_model import Isobox
from libisobox.isoboxmanagement import getboxbyname
from libisobox.checkdependencies import checkdependencies
from libisobox.purge import purge
from libisobox.cleanup_processes import cleanup_processes
from libisobox.signalhandler import Signalhandler
from libisobox.create_isobox import create_isobox
from libisobox.list_isoboxes import list_isoboxes
from libisobox.remove_isobox import remove_isobox
from libisobox.move_isobox import move_isobox


def create_required_files():  # Create files used by isobox like config files if they dont exist
    if not os.path.exists("/var/lib/isobox"):
        os.mkdir("/var/lib/isobox")
    if not os.path.exists("/var/lib/isobox/isoboxes.json"):
        with open("/var/lib/isobox/isoboxes.json", "w") as f:
            json.dump([], f)
    if not os.path.exists("/var/lib/isobox/mounts"):
        os.mkdir("/var/lib/isobox/mounts")


def parse_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(help="Command", dest="command")
    subparsers.required = True

    createparser = subparsers.add_parser("create", help="Create new isobox")
    createparser.add_argument("name")
    createparser.add_argument("isopath")
    createparser.add_argument(
        "-sfs",
        "--sfspath",
        type=str,
        default=None,
        help="Path of the iso's squashfs relative to the iso's root",
    )

    removeparser = subparsers.add_parser("remove", help="Remove existing isobox")
    removeparser.add_argument("name")

    rmallparser = subparsers.add_parser(
        "purge",
        help="Remove all existing isoboxes and data related to isobox, usually to run before uninstalling",
    )

    infoparser = subparsers.add_parser("info", help="Print info about an isobox")
    infoparser.add_argument("name")

    mvparser = subparsers.add_parser(
        "move", help="Move an isobox to another directory on your system."
    )
    mvparser.add_argument("name")
    mvparser.add_argument("target")

    createparser.add_argument("-mountpoint", type=str, default=None)

    shellparser = subparsers.add_parser(
        "shell", help="Gain a shell into an existing isobox"
    )
    shellparser.add_argument("name")
    shellparser.add_argument(
        "-u",
        "--user",
        type=str,
        default="root",
        help="The user inside the isobox to start as",
    )

    runparser = subparsers.add_parser("run", help="Run program in an isobox")
    runparser.add_argument("name")
    runparser.add_argument("program")
    runparser.add_argument(
        "-u",
        "--user",
        type=str,
        default="root",
        help="The user inside the isobox to start as",
    )

    subparsers.add_parser("ls", help="List all isoboxes")

    guiparser = subparsers.add_parser("gui", help="Run existing isobox and start gui")
    guiparser.add_argument("name")
    guiparser.add_argument("-tty", type=str, default="3", help="The tty to startx into")
    guiparser.add_argument(
        "-display",
        type=str,
        default="2",
        help="The display to start the Xorg server on",
    )
    guiparser.add_argument(
        "-u",
        "--user",
        type=str,
        default="root",
        help="The user inside the isobox to start as",
    )

    return parser.parse_args()


def main(args):

    create_required_files()

    if args.command == "create":
        create_isobox(args)

    elif args.command == "shell":
        currentbox = getboxbyname(args.name)

        mount_filesystems(currentbox["mountpoint"])
        chroot(currentbox["mountpoint"], args.user)

    elif args.command == "gui":
        currentbox = getboxbyname(args.name)

        mount_filesystems(currentbox["mountpoint"])
        chroot_gui(currentbox["mountpoint"], args.tty, args.display, args.user)
    elif args.command in ("ls", "list"):
        list_isoboxes()

    elif args.command == "remove":
        remove_isobox(args)

    elif args.command == "move":
        move_isobox(args)

    elif args.command == "info":
        currentbox = getboxbyname(args.name)
        print(f"Info about box {currentbox['name']}:\n")
        print(f"Name: {currentbox['name']}")
        print(f"Mountpoint: {currentbox['mountpoint']}")
        print(f"Note: {currentbox['note']}")
    elif args.command == "purge":
        answer = input(
            "Are you sure you want to delete all data related to isobox? [Y/n]: "
        ).lower()
        if answer in ("y", "yes"):
            purge()
        else:
            sys.exit("Aborting")
    elif args.command == "run":
        currentbox = getboxbyname(args.name)

        mount_filesystems(currentbox["mountpoint"])
        chroot_run(currentbox["mountpoint"], args.user, args.program)


if __name__ == "__main__":
    args = parse_args()

    if (
        os.getuid() != 0
    ):  # Quit if user is not root as root is pretty much required for this to run properly
        sys.exit("Requires root")
    checkdependencies()  # Check if unsquashfs is installed basically

    currentmountpoint = None
    signalhandler = None
    # Set the currentmountpoint here so we can unmount mounted filesystems  here if something goes south
    if args.command in ("shell", "gui", "run") and getboxbyname(args.name):
        currentmountpoint = getboxbyname(args.name)["mountpoint"]
        signalhandler = Signalhandler(currentmountpoint)

    try:
        main(args)
    except Exception as e:
        print(traceback.format_exc())

    finally:
        if currentmountpoint:
            if signalhandler.dead:
                sys.exit(0)
            else:
                umount_filesystems(currentmountpoint)
                cleanup_processes(currentmountpoint)