AdvLabDB/manage.py

#!/usr/bin/python3

import secrets
import subprocess  # nosec 404
from getpass import getpass
from pathlib import Path
from shutil import copytree, rmtree

import click
from email_validator import validate_email
from flask_admin import __file__ as flask_admin_path
from flask_security import admin_change_password, hash_password
from sqlalchemy import select

from advlabdb import app, db, settings, user_datastore
from advlabdb.model_independent_funs import randomPassword
from advlabdb.models import MAX_YEAR, MIN_YEAR, Admin, Semester, User


def run(command: str, **kwargs):
    return subprocess.run(command, shell=True, **kwargs)  # nosec B602


def box(message: str):
    click.echo()
    click.echo(click.style(f" {message} ", bg="white", fg="black"))
    click.echo()


# Class to validate email in click.prompt
class EmailParamType(click.ParamType):
    def convert(self, value, param, ctx):
        try:
            return validate_email(value).email
        except Exception:
            self.fail(f"{value} is not a valid email!", param, ctx)


@click.group(
    help="Command line tool to manage AdvLabDB.",
)
def cli():
    pass


@cli.group(
    short_help="Setup commands.",
    help="Commands used to setup an AdvLabDB instance.",
)
def setup():
    pass


@setup.command(
    short_help="Generate required secrets.",
    help="Generate the file secrets.ini if it does not already exist.",
)
def generate_secrets():
    file = Path("secrets.ini")

    if file.is_file():
        click.echo(f"Skipping secrets generation because the secrets file does already exist at {file}.")
        return

    with open(file, "w") as f:
        f.write("[Secrets]\n")

        key = secrets.SystemRandom().getrandbits(128)
        f.write(f"SECRET_KEY = {key}\n")

        salt = secrets.token_hex()
        f.write(f"SECURITY_PASSWORD_SALT = {salt}\n")


@setup.command(
    short_help="Initialize the database.",
    help="Initialize the database if it does not already exist.",
)
def initialize_database():
    db_file = Path(settings["SQLITE_DB_PATH"])
    if db_file.is_file():
        click.echo(f"Skipping database initialization because the database does already exist at {db_file}.")
        return

    click.echo("\nThis script should only be used to initialize the database after setting up a server")
    click.echo(click.style("The old database will be DELETED and a new database will be created!", bg="red"))

    if not click.confirm(click.style("Are you sure that you want to continue?", fg="red"), default=False):
        click.echo(click.style("Aborted!", fg="yellow"))
        return

    with app.app_context():
        with db.session.begin():
            # Delete old database
            db.drop_all()
            # Create new database
            db.create_all()

            semester_label = click.prompt(
                "Enter the label of the current semester",
                type=click.Choice(("SS", "WS")),
            )
            semester_year = click.prompt(
                f"Enter the year of the current semester (between {MIN_YEAR} and {MAX_YEAR})",
                type=click.IntRange(MIN_YEAR, MAX_YEAR),
            )

            semester = Semester(label=semester_label, year=semester_year)

            db.session.add(semester)

            adminRole = user_datastore.create_role(name="admin")
            user_datastore.create_role(name="assistant")

            box("The first admin account will be created now.")

            admin_email = click.prompt(
                "Enter the admin's email address",
                type=EmailParamType(),
            )

            admin_first_name = click.prompt("Enter the admin's first name")
            admin_last_name = click.prompt("Enter the admin's last name")
            admin_phone_number = click.prompt(
                "Enter the admin's phone number (optional)", default="", show_default=False
            )
            admin_mobile_phone_number = click.prompt(
                "Enter the admin's mobile phone number (optional)", default="", show_default=False
            )
            admin_building = click.prompt("Enter the admin's building (optional)", default="", show_default=False)
            admin_room = click.prompt("Enter the admin's room (optional)", default="", show_default=False)

            admin_password = randomPassword()
            admin_hashed_password = hash_password(admin_password)

            admin_user = user_datastore.create_user(
                email=admin_email,
                password=admin_hashed_password,
                roles=[adminRole],
                first_name=admin_first_name.strip(),
                last_name=admin_last_name.strip(),
                phone_number=admin_phone_number.strip() or None,
                mobile_phone_number=admin_mobile_phone_number.strip() or None,
                building=admin_building.strip() or None,
                room=admin_room.strip() or None,
                active_semester=semester,
            )

            admin = Admin(user=admin_user)

            db.session.add(admin)

    box(f"Admin password: {admin_password}")

    click.echo(click.style("Done database initialization!", fg="green"))


@cli.group(
    short_help="Maintainance commands.",
    help="Commands used to maintain an AdvLabDB instance.",
)
def maintain():
    pass


@maintain.command(
    short_help="Reset the password of an admin.",
    help="Reset the password of a chosen active admin. A random password will be generated. If no admins are active, a chosen admin will be activated.",
)
def reset_admin_password():
    click.echo("This script will generate a new random password for a chosen admin.\n")

    with app.app_context():
        with db.session.begin():
            admins = db.session.execute(select(Admin).join(User).where(User.active == True)).scalars().all()
            activate_user = False

            if len(admins) == 0:
                click.echo("There is no admin with an active user. The user of the chosen admin will be activated.")
                admins = db.session.execute(select(Admin)).scalars().all()
                activate_user = True

            num_admins = len(admins)

            prompt = "Admins:\n"
            for ind, admin in enumerate(admins):
                user = admin.user
                prompt += f"[{ind}] {user.first_name} {user.last_name}: {user.email}\n"
            prompt += f"Enter number [0-{num_admins - 1}]"

            admin_index = click.prompt(
                prompt,
                type=click.IntRange(0, num_admins - 1),
            )

            chosen_admin_user = admins[admin_index].user

            new_password = randomPassword()

            admin_change_password(
                chosen_admin_user, new_password, notify=False
            )  # Password is automatically hashed with this function

            if activate_user:
                chosen_admin_user.active = True

    box(f"New password: {new_password}")

    click.echo(click.style("Done!", fg="green"))


@maintain.command(
    short_help="Copy admin templates",
    help="Copy the templates from the Flask-Admin package. This is only needed if the templates should be updated to a new version after a new release of Flask-Admin.",
)
def copy_admin_templates():
    src = Path(flask_admin_path).parent / "templates/bootstrap4/admin"
    if not src.is_dir():
        click.echo(click.style(f"Templates could not be found at {src}", fg="red"))
        return

    dist = Path("advlabdb/templates/admin")
    if dist.is_dir():
        while True:
            ans = input().lower()
            if ans == "s":
                print("Process stopped!")
                return False
            elif ans == "o":
                break
        if not click.confirm(
            click.style(f"The directory {dist} already exists! Do you want to overwrite it?", fg="yellow")
        ):
            return

        rmtree(dist)
        click.echo(click.style("Old templates deleted!", fg="yellow"))

    copytree(src, dist)
    click.echo(click.style(f"Copied {src} -> {dist}", fg="green"))

    click.echo(
        click.style(
            f"""
_________
| WARNING
| -------
| You might have to edit the file {dist}/base.html
| by adding nav in the following way:
| This line:\t<ul class="navbar-nav mr-auto">
| Becomes:\t<ul class="nav navbar-nav mr-auto">
|
| This will prevent the navigation bar from expanding
| such that some elements can not be seen.
| Refer to this pull request:
| https://github.com/flask-admin/flask-admin/pull/2233
|
| If the above pull request is merged and flask-admin
| is on a new release after the merge,
| then this step is not needed.
_________
        """,
            fg="yellow",
        )
    )


if __name__ == "__main__":
    cli()