user

#!/usr/bin/python3

import os
import yaml
import click
import subprocess

from urllib.request import urlopen
from user_modules import modules


class colors:
    reset = '\033[0m'
    bold = '\033[01m'
    disable = '\033[02m'
    underline = '\033[04m'
    reverse = '\033[07m'
    strikethrough = '\033[09m'
    invisible = '\033[08m'

    class fg:
        black = '\033[30m'
        red = '\033[31m'
        green = '\033[32m'
        orange = '\033[33m'
        blue = '\033[34m'
        purple = '\033[35m'
        cyan = '\033[36m'
        lightgrey = '\033[37m'
        darkgrey = '\033[90m'
        lightred = '\033[91m'
        lightgreen = '\033[92m'
        yellow = '\033[93m'
        lightblue = '\033[94m'
        pink = '\033[95m'
        lightcyan = '\033[96m'

        rainbow = [lightred, orange, yellow,
                   lightgreen, lightcyan, blue, purple]
        seq = 0

        def random(self):
            if self.seq == 7:
                self.seq = 0
            self.seq += 1
            return self.rainbow[self.seq - 1]

        def clear_seq(self):
            self.seq = 0

    class bg:
        black = '\033[40m'
        red = '\033[41m'
        green = '\033[42m'
        orange = '\033[43m'
        blue = '\033[44m'
        purple = '\033[45m'
        cyan = '\033[46m'
        lightgrey = '\033[47m'


fg = colors.fg()


def info(msg):
    print(colors.bold + fg.cyan + '[INFO] ' +
          colors.reset + msg + colors.reset)


def print_list(msg):
    print(colors.bold + fg.random() + '[LIST] ' +
          colors.reset + msg + colors.reset)


def modrun(msg):
    print(colors.bold + fg.green + '[MODRUN] ' +
          colors.reset + msg + colors.reset)


def container_msg(msg):
    print(colors.bold + fg.purple + '[CONTAINER] ' +
          colors.reset + msg + colors.reset)


def association_msg(msg):
    print(colors.bold + fg.random() + '[ASSOCIATION] ' +
          colors.reset + msg + colors.reset)


def warn(warning):
    print(colors.bold + fg.yellow + '[WARNING] ' +
          colors.reset + warning + colors.reset)


def error(err):
    print(colors.bold + fg.red + '[ERROR] ' +
          colors.reset + err + colors.reset)


def proceed():
    print(colors.bold + fg.red + '[QUESTION] ' +
          colors.reset + 'would you like to proceed?' + colors.reset)
    info(f'(press {colors.bold}ENTER{colors.reset} to proceed, or {colors.bold}^C{colors.reset}/{colors.bold}^D{colors.reset} to cancel)')
    input()


@click.group("cli")
def cli():
    """Manage user operations using the user utility on blendOS."""


def main():
    cli(prog_name="user")


@cli.command("cadre")
@click.argument('path', type=click.Path(exists=True))
def apply_cadre(path):
    '''
    Apply a cadre (TOML configuration file)
    '''

    config = {}

    with open(path) as f:
        config = yaml.safe_load(f)

    if type(config.get('associations')) == dict:
        modrun('invoking all modules')

        files = {}
        commands = []

        for module in config['modules'].keys():
            try:
                modules[module]
                if config['modules'][module].get('enabled') == True:
                    print()
                    modrun(
                        f'invoking module {colors.bold}{fg.green}{module}{colors.reset}')
                    mod_return = modules[module].call(
                        config['modules'][module])
                    if type(mod_return.get('files')) == dict:
                        files.update(mod_return['files'])

                    if type(mod_return.get('commands')) == list:
                        commands += mod_return['commands']

                    modrun(f'{module} completed')
            except KeyError:
                print()
                warn(f'could not find module {module}, skipping')
                exit(1)

        print()
        info('all modules have been executed')
        if len(files) != 0:
            print()
            info('the following files will be created:')
            for i, file in enumerate(files.keys()):
                overwriting = ''
                if os.path.isfile(os.path.expanduser(f'~/{file}')):
                    overwriting = f' {fg.orange}[OVERWRITING]'
                print_list(
                    f' {colors.bold}{i+1}. {file}{overwriting}{colors.reset}')
            fg.clear_seq()

        if len(commands) != 0:
            print()
            info('the following commands will be run:')
            for i, command in enumerate(commands):
                print_list(f' {colors.bold}{i+1}. {command}{colors.reset}')
            fg.clear_seq()

        print()
        proceed()

        for f in files.keys():
            subprocess.run(['mkdir', '-p', os.path.dirname(f)])
            try:
                print_list(f'writing file {colors.bold}{f}{colors.reset}')
                with open(os.expanduser(f'~/{f}'), 'w') as f_o:
                    f_o.write(files[f])
            except IsADirectoryError:
                warn(f'{f} is a directory, skipping')
        fg.clear_seq()

        print()

        for c in commands:
            print_list(f'running command {colors.bold}{c}{colors.reset}')
            subprocess.run(c)
        fg.clear_seq()

        if len(commands) != 0:
            print()

        modrun(
            'successfully run all modules (you might have to relogin for all changes to take effect)')
        print()

    if type(config.get('containers')) == dict:
        container_msg('creating all containers')
        print()

        for container in config['containers'].keys():
            distro = 'arch'
            packages = []
            commands = []
            if type(config['containers'][container]) == str:
                distro = config['containers'][container]
            elif type(config['containers'][container]) == dict:
                if type(config['containers'][container].get('distro')) == str:
                    distro = config['containers'][container].get('distro')
                if type(config['containers'][container].get('packages')) == list:
                    packages = config['containers'][container].get('packages')
                if type(config['containers'][container].get('commands')) == list:
                    for cmd in config['containers'][container].get('commands'):
                        if type(cmd) == list:
                            commands.append(cmd)
                        elif type(cmd) == str:
                            commands.append(['bash', '-c', cmd])
            creation_env = os.environ.copy()
            creation_env['BLEND_NO_CHECK'] = 'true'
            container_msg(
                f'{colors.bold}{container}{colors.reset} is being created')
            if subprocess.run(['podman', 'container', 'exists', container], env=creation_env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0:
                if distro not in ('arch', 'almalinux-9', 'crystal-linux', 'debian', 'fedora-38', 'kali-linux', 'neurodebian-bookworm', 'rocky-linux', 'ubuntu-22.04', 'ubuntu-23.04'):
                    warn(
                        f'distro {colors.bold}{distro}{colors.reset} not supported, skipping')
                    continue
                if subprocess.run(['blend', 'create-container', '-cn', container, '-d', distro], env=creation_env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0:
                    warn(
                        f'creating {colors.bold}{container}{colors.reset} failed, skipping')
                    continue
                else:
                    container_msg(
                        f'created {colors.bold}{container}{colors.reset} successfully')
            else:
                warn(f'{colors.bold}{container}{colors.reset} already exists')

            print()

            if len(packages) != 0:
                info(
                    f'installing the following packages in container {colors.bold}{container}{colors.reset}: {colors.bold}{f"{colors.reset}, {colors.bold}".join(packages)}{colors.reset}')
                if distro in ('arch', 'crystal-linux'):
                    subprocess.run(
                        [f'sudo.{container}', 'pacman', '-Syu', *packages])
                elif distro in ('ubuntu-22.04', 'ubuntu-23.04', 'debian', 'neurodebian-bookworm', 'kali-linux'):
                    subprocess.run([f'sudo.{container}', 'apt-get', 'update'])
                    subprocess.run(
                        [f'sudo.{container}', 'apt-get', 'install', '-y', *packages])
                elif distro in ('almalinux-9', 'fedora-38', 'rocky-linux'):
                    subprocess.run(
                        [f'sudo.{container}', 'dnf', 'check-update'])
                    subprocess.run(
                        [f'sudo.{container}', 'dnf', 'install', '-y', *packages])
                info(
                    f'installed packages for container {colors.bold}{container}{colors.reset}')
                print()

            if len(commands) != 0:
                info(
                    f'running custom commands for {colors.bold}{container}{colors.reset}')
                for cmd in commands:
                    creation_env = os.environ.copy()
                    creation_env['BLEND_NO_CHECK'] = 'true'
                    subprocess.run(['blend', 'enter', '-cn', container, '--', *cmd], env=creation_env)
                info(
                    f'finished running all commands for {colors.bold}{container}{colors.reset}')
                print()

        container_msg('created all containers')
        print()

    if type(config.get('associations')) == dict:
        info('creating all associations')

        if len(config['associations'].keys()) != 0:
            print()

        fg.clear_seq()

        for association in config['associations'].keys():
            container = config['associations'][association]
            if not os.path.exists(os.path.expanduser(f'~/.local/bin/blend_bin/{association}.{container}')):
                warn(f'{colors.bold}{association}.{container}{colors.reset} does not exist, skipping creation of {colors.bold}{association}{colors.reset}')
                continue
            if os.path.isfile(os.path.expanduser('~/.local/bin/blend_bin/.associations')):
                subprocess.run(['sed', '-i', f's/^{association}\\x0//g', os.path.expanduser(
                    '~/.local/bin/blend_bin/.associations')])
            with open(os.path.expanduser('~/.local/bin/blend_bin/.associations'), 'a+') as f:
                f.write(f'{association}\0{container}\n')
            subprocess.run(['ln', '-sf', f'{association}.{container}',
                           os.path.expanduser(f'~/.local/bin/blend_bin/{association}')])
            association_msg(
                f'created {colors.bold}{association} -> {container}{colors.reset}')

        print()
        info('created all associations')
        print()

        print(colors.bold + fg.green + '[COMPLETE] ' + colors.reset +
              'the training of the cadre is now complete' + colors.reset)


@cli.command("associate")
@click.argument('association')
@click.argument('container')
def associate_binary(association, container):
    '''
    Create an association (for example, apt -> ubuntu)
    '''

    if not os.path.exists(os.path.expanduser(f'~/.local/bin/blend_bin/{association}.{container}')):
        error(f'{colors.bold}{association}.{container}{colors.reset} does not exist')
        exit()
    if os.path.isfile(os.path.expanduser('~/.local/bin/blend_bin/.associations')):
        subprocess.run(['sed', '-i', f's/^{association}\\x0//g',
                       os.path.expanduser('~/.local/bin/blend_bin/.associations')])
    with open(os.path.expanduser('~/.local/bin/blend_bin/.associations'), 'a+') as f:
        f.write(f'{association}\0{container}\n')
    _exists = os.path.exists(os.path.expanduser(
        f'~/.local/bin/blend_bin/{association}'))
    subprocess.run(['ln', '-sf', f'{association}.{container}',
                   os.path.expanduser(f'~/.local/bin/blend_bin/{association}')])
    association_msg(('modified' if _exists else 'created') +
                    f' {colors.bold}{association} -> {container}{colors.reset}')


@cli.command("dissociate")
@click.argument('association')
def associate_binary(association):
    '''
    Remove an association
    '''

    if not os.path.exists(os.path.expanduser(f'~/.local/bin/blend_bin/{association}')):
        error(f'{colors.bold}{association}{colors.reset} does not exist')
        exit()
    if os.path.isfile(os.path.expanduser('~/.local/bin/blend_bin/.associations')):
        subprocess.run(['sed', '-i', f's/^{association}\\x0//g',
                       os.path.expanduser('~/.local/bin/blend_bin/.associations')])
    subprocess.run(
        ['rm', '-f', os.path.expanduser(f'~/.local/bin/blend_bin/{association}')])
    association_msg(f'dissociated {colors.bold}{association}')


@cli.command("create-container")
@click.argument('container_name')
@click.argument('distro', default='arch')
def create_container(container_name, distro):
    '''
    Create a container
    '''
    if distro not in ('arch', 'almalinux-9', 'crystal-linux', 'debian', 'fedora-38', 'kali-linux', 'neurodebian-bookworm', 'rocky-linux', 'ubuntu-22.04', 'ubuntu-23.04'):
        error(
            f'distro {colors.bold}{distro}{colors.reset} not supported')
    if subprocess.run(['podman', 'container', 'exists', container_name], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0:
        error(f'container {colors.bold}{container_name}{colors.reset} already exists')
        exit(1)
    subprocess.run(['blend', 'create-container', '-cn', container_name, '-d', distro])


@cli.command("delete-container")
@click.argument('container')
def delete_container(container):
    '''
    Delete a container
    '''
    if subprocess.run(['podman', 'container', 'exists', container], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0:
        error(f'container {colors.bold}{container}{colors.reset} does not exist')
        exit(1)
    subprocess.run(['blend', 'remove-container', container])


@cli.command("shell")
@click.argument('container')
def shell(container):
    '''
    Enter a shell inside a container
    '''
    if subprocess.run(['podman', 'container', 'exists', container], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0:
        error(f'container {colors.bold}{container}{colors.reset} does not exist')
        exit(1)
    creation_env = os.environ.copy()
    creation_env['BLEND_NO_CHECK'] = 'true'
    subprocess.run(['blend', 'enter', '-cn', container], env=creation_env)


@cli.command("exec")
@click.argument('container')
@click.argument('cmds', nargs=-1, required=True)
def exec_c(container, cmds):
    '''
    Run a command inside a container
    '''
    if subprocess.run(['podman', 'container', 'exists', container], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0:
        error(f'container {colors.bold}{container}{colors.reset} does not exist')
        exit(1)
    creation_env = os.environ.copy()
    creation_env['BLEND_NO_CHECK'] = 'true'
    subprocess.run(['blend', 'enter', '-cn', container, '--', *cmds], env=creation_env)


@cli.command("install")
@click.argument('container')
@click.argument('pkgs', nargs=-1, required=True)
def install_c(container, pkgs):
    '''
    Install a package inside a container
    '''
    if os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/apt.{container}')):
        subprocess.run([f'sudo.{container}', 'apt', 'update'])
        subprocess.run([f'sudo.{container}', 'apt', 'install', *pkgs])
    elif os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/dnf.{container}')):
        subprocess.run([f'sudo.{container}', 'dnf', 'install', *pkgs])
    elif os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/pacman.{container}')):
        subprocess.run([f'sudo.{container}', 'pacman', '-Syu', *pkgs])
    else:
        error(f'container {colors.bold}{container}{colors.reset} does not exist')
        exit(1)


@cli.command("remove")
@click.argument('container')
@click.argument('pkgs', nargs=-1, required=True)
def remove_c(container, pkgs):
    '''
    Remove a package inside a container
    '''
    if os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/apt.{container}')):
        subprocess.run([f'sudo.{container}', 'apt', 'purge', *pkgs])
    elif os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/dnf.{container}')):
        subprocess.run([f'sudo.{container}', 'dnf', 'remove', *pkgs])
    elif os.path.isfile(os.path.expanduser(f'~/.local/bin/blend_bin/pacman.{container}')):
        subprocess.run([f'sudo.{container}', 'pacman', '-Rcns', *pkgs])
    else:
        error(f'container {colors.bold}{container}{colors.reset} does not exist')
        exit(1)


if __name__ == '__main__':
    main()