audius-cli

#!/usr/bin/env python3

import json
import logging
import os
import pathlib
import random
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request

import click
import crontab
import dotenv
import psutil

from web3 import Web3

from web3.providers import BaseProvider, HTTPProvider

GB = 1024**3

RECOMMENDED_CPU_COUNT = 8
RECOMMENDED_MEMORY = 16 * GB
RECOMMENDED_STORAGE = {
    "creator-node": 2048 * GB,
    "discovery-provider": 256 * GB,
    "identity-service": 256 * GB,
}

SERVICE_PORTS = {
    "discovery-provider": "5000",
    "discovery-provider-notifications": "6000",
    "creator-node": "4000",
    "identity-service": "7000",
}

SERVICES = (
    "creator-node",
    "discovery-provider",
    "identity-service",
)

# if type == discovery-provider and is not in this list, run core only
USE_LEGACY_DISCOVERY_PROVIDER_PROFILE = (
    ".audius.co",
    ".creatorseed.com",
    "dn1.monophonic.digital",
    ".figment.io",
    ".tikilabs.com",
)

REGISTERED_PLUGINS = "REGISTERED_PLUGINS"
DB_URL = "audius_db_url"

service_type = click.Choice(SERVICES)
network_type = click.Choice(["prod", "stage", "dev"])
container_type = click.Choice(["backend", "cache", "db"])


def run(cmd, **kwargs):
    print(subprocess.list2cmdline(cmd))
    return subprocess.run(cmd, **kwargs)


def get_network(ctx, service):
    """Returns the network name for the given service."""
    return dotenv.dotenv_values(ctx.obj["manifests_path"] / service / ".env").get(
        "NETWORK", "prod"
    )


def get_override_path(ctx, service):
    """Returns path to override.env for the given service."""
    return pathlib.Path(
        dotenv.dotenv_values(ctx.obj["manifests_path"] / service / ".env").get(
            "OVERRIDE_PATH", ctx.obj["manifests_path"] / service / "override.env"
        )
    )


def lock(ctx, group):
    lockfile = pathlib.Path(f"~/.local/share/audius-cli/{group}.lock").expanduser()
    lockfile.parent.mkdir(parents=True, exist_ok=True)
    lockfile.touch(exist_ok=True)

    pid, retries = lockfile.read_text(), 0
    while pid.isdigit() and int(pid) != os.getpid() and psutil.pid_exists(int(pid)):
        if retries == 600:  # 10 mins
            click.secho("Giving up on lock")
            sys.exit(1)
        if retries % 15 == 0:
            click.secho(f"Waiting for lock (held by pid {pid})", color="gray")
        time.sleep(1)
        pid = lockfile.read_text()
        retries += 1

    lockfile.write_text(str(os.getpid()))


def set_automatic_env(ctx):
    with crontab.CronTab(user=True) as cron:
        auto_upgrade = (
            "true" if any(cron.find_comment("audius-cli auto-upgrade")) else "false"
        )

    git_hash = (
        run(
            ["git", "rev-parse", "HEAD"],
            capture_output=True,
            cwd=ctx.obj["manifests_path"],
        )
        .stdout.decode()
        .rstrip()
    )

    # give application awareness of host level auto upgrade cron job

    ctx.invoke(
        set_config,
        service="creator-node",
        key="autoUpgradeEnabled",
        value=auto_upgrade,
    )

    ctx.invoke(
        set_config,
        service="discovery-provider",
        key="audius_auto_upgrade_enabled",
        value=auto_upgrade,
    )

    ctx.invoke(
        set_config,
        service="discovery-provider",
        key="AUDIUS_DOCKER_COMPOSE_GIT_SHA",
        value=git_hash,
    )

    ctx.invoke(
        set_config,
        service="creator-node",
        key="AUDIUS_DOCKER_COMPOSE_GIT_SHA",
        value=git_hash,
    )


def prune():
    disk_usage_percent = psutil.disk_usage("/").percent
    if disk_usage_percent > 90:
        run(
            ["docker", "system", "prune", "--all", "--force"],
        )
    else:
        run(
            ["docker", "system", "prune", "--all", "--force", "--filter", "until=48h"],
        )


@click.group()
@click.pass_context
def cli(ctx):
    """A tool for managing audius services"""
    ctx.ensure_object(dict)
    ctx.obj["manifests_path"] = pathlib.Path(
        os.getenv("MANIFESTS_PATH", os.path.dirname(os.path.realpath(__file__)))
    )

    logging.basicConfig(
        filename=ctx.obj["manifests_path"] / "audius-cli.log",
        level=logging.INFO,
        format="%(asctime)s:%(levelname)s:%(message)s",
    )


@cli.command()
@click.argument("service", type=service_type)
@click.pass_context
def check_config(ctx, service):
    """Check the config for a service"""
    env = ctx.obj["manifests_path"] / service / f"{get_network(ctx, service)}.env"
    override_env = get_override_path(ctx, service)

    env_data = dotenv.dotenv_values(env)
    override_env_data = dotenv.dotenv_values(override_env)

    unset = False
    for key, value in env_data.items():
        if override_env_data.get(key, value) == "":
            unset = True
            click.secho(f"{key} is not set", fg="red")

    if unset:
        sys.exit(1)
    else:
        click.secho("All keys are set", fg="green")


@cli.command()
@click.argument("service", type=service_type)
@click.pass_context
def health_check(ctx, service):
    """Check the health of a service"""

    path = ctx.obj["manifests_path"] / service
    container = "backend"
    if service == "discovery-provider-notifications":
        path = ctx.obj["manifests_path"] / "discovery-provider"
        container = "notifications"
    elif service == "creator-node":
        container = "audiusd"

    proc = run(
        [
            "docker",
            "compose",
            "--project-directory",
            path,
            "ps",
            "-q",
            container,
        ],
        capture_output=True,
    )

    if proc.returncode:
        click.secho("Service is not running", fg="yellow")
        sys.exit(1)

    try:
        response = json.load(
            urllib.request.urlopen(
                f"http://localhost:{SERVICE_PORTS[service]}/health_check"
            )
        )

        click.secho("Response:", bold=True)
        click.echo(json.dumps(response, indent=2, sort_keys=True))

        partial = False
        if service == "creator-node":
            healthy = response["data"]["healthy"]
        elif service == "discovery-provider":
            healthy = "block_difference" in response["data"]
            block_diff = response["data"]["block_difference"]
            max_block_diff = response["data"]["maximum_healthy_block_difference"]
            partial = block_diff > max_block_diff
            if block_diff > max_block_diff:
                click.secho(
                    f"Block difference ({block_diff}) is greater than maximum healthy block difference ({max_block_diff})",
                    fg="yellow",
                )
        elif service == "identity-service":
            healthy = response["healthy"]
        elif service == "discovery-provider-notifications":
            healthy = response["healthy"]

        if healthy and partial:
            click.secho("Service is partially healthy", fg="yellow")
            sys.exit(2)
        elif healthy:
            click.secho("Service is healthy", fg="green")
            sys.exit(0)
        else:
            click.secho("Service is not healthy", fg="red")
            sys.exit(1)
    except (
        ConnectionError,
        ConnectionRefusedError,
        urllib.error.HTTPError,
        urllib.error.URLError,
        json.JSONDecodeError,
    ):
        click.secho("Service is not healthy", fg="red")
        sys.exit(1)


@cli.command()
@click.argument("service", type=click.Choice(('creator-node', 'discovery-provider')))
@click.pass_context
def min_version(ctx, service):
    """Check the minimum required version for a service"""
    env = ctx.obj["manifests_path"] / service / f"{get_network(ctx, service)}.env"
    override_env = get_override_path(ctx, service)

    env_data = dotenv.dotenv_values(env)
    override_env_data = dotenv.dotenv_values(override_env)

    if service == "creator-node":
        service = "content-node"
    elif service == "discovery-provider":
        service = "discovery-node"

    rpc_env_var = "audius_web3_eth_provider_url" if service == "discovery-node" else "ethProviderUrl"
    rpc_url = override_env_data.get(rpc_env_var, env_data.get(rpc_env_var))
    if not rpc_url:
        click.secho(f"{rpc_env_var} env var is not set", fg="red")
        sys.exit(1)
    
    eth_registry_address_env_var = "audius_eth_contracts_registry" if service == "discovery-node" else "ethRegistryAddress"
    eth_registry_address_str = override_env_data.get(eth_registry_address_env_var, env_data.get(eth_registry_address_env_var))
    if not eth_registry_address_str:
        click.secho(f"{eth_registry_address_env_var} env var is not set", fg="red")
        sys.exit(1)

    eth_registry_address = Web3.to_checksum_address(eth_registry_address_str)
    eth_web3 = get_eth_web3(rpc_url)
    eth_abi_values = load_eth_abi_values(os.path.join(ctx.obj["manifests_path"], "eth-contracts", "ABIs"))
    eth_registry_instance = eth_web3.eth.contract(
        address=eth_registry_address, abi=eth_abi_values["Registry"]["abi"]
    )
    contract_address = eth_registry_instance.functions.getContract(
        "ServiceTypeManagerProxy".encode("utf-8")
    ).call()
    contract_instance = eth_web3.eth.contract(
        address=contract_address, abi=eth_abi_values["ServiceTypeManagerProxy"]["abi"]
    )
    min_version = contract_instance.functions.getCurrentVersion(
        service.encode("utf-8")
    ).call()
    click.secho(f"Minimum required version for {service} is {min_version.decode('utf-8')}", fg="green")
    sys.exit(0)


def load_eth_abi_values(abi_dir):
    json_files = os.listdir(abi_dir)
    loaded_abi_values = {}
    for contract_json_file in json_files:
        fullPath = os.path.join(abi_dir, contract_json_file)
        with open(fullPath) as f:
            data = json.load(f)
            loaded_abi_values[data["contractName"]] = data
    return loaded_abi_values


def get_eth_web3(rpc_url):
    provider = MultiProvider(rpc_url)
    for p in provider.providers:
        p.middlewares.clear()
    # Remove the default JSON-RPC retry middleware
    # as it correctly cannot handle eth_getLogs block range
    # throttle down.
    # See https://web3py.readthedocs.io/en/stable/examples.html
    eth_web3 = Web3(provider)
    eth_web3.strict_bytes_type_checking = False
    return eth_web3


class MultiProvider(BaseProvider):
    """
    Implements a custom web3 provider

    ref: https://web3py.readthedocs.io/en/stable/internals.html#writing-your-own-provider
    """

    def __init__(self, providers):
        self.providers = [HTTPProvider(provider) for provider in providers.split(",")]

    def make_request(self, method, params):
        for provider in random.sample(self.providers, k=len(self.providers)):
            try:
                return provider.make_request(method, params)
            except Exception:
                continue
        raise Exception("All requests failed")

    def isConnected(self):
        return any(provider.isConnected() for provider in self.providers)

    def __str__(self):
        return f"MultiProvider({self.providers})"

@cli.command()
@click.option("-y", "--yes", is_flag=True)
@click.option("--seed", is_flag=True, help="Seed the discovery-provider database")
@click.option(
    "--auto-seed", 
    is_flag=True, 
    help="Seed the discovery-provider database if it's empty (e.g. for new nodes)",
)
@click.option("--chain", is_flag=True)
@click.argument("service", type=service_type)
@click.pass_context
def launch(ctx, service, seed, auto_seed, chain, yes):
    """Launch the service"""
    set_automatic_env(ctx)

    lock(ctx, "docker")

    log_file = ctx.obj["manifests_path"] / "auto-upgrade.log"
    try:
        # clear the previous auto-upgrade logs
        open(log_file, "w").close()
    except Exception as e:
        print(f"Error while clearing the log file: {e}")

    try:
        ctx.invoke(check_config, service=service)
    except SystemExit:
        pass
    total_memory = psutil.virtual_memory().total
    cpu_count = psutil.cpu_count()
    total_storage = shutil.disk_usage("/var/k8s").total

    click.echo(f"CPUs:\t{cpu_count}\t(required: {RECOMMENDED_CPU_COUNT})")
    click.echo(
        f"Memory:\t{total_memory // GB}GB\t(required: {RECOMMENDED_MEMORY // GB}GB)"
    )
    click.echo(
        f"Storage:\t{total_storage // GB}GB\t(required: {RECOMMENDED_STORAGE[service] // GB}GB)"
    )

    if (
        cpu_count < RECOMMENDED_CPU_COUNT
        or total_memory < RECOMMENDED_MEMORY
        or total_storage < RECOMMENDED_STORAGE[service]
    ):
        click.secho("System does not meet requirements", fg="red")
    else:
        click.secho("System meets requirements", fg="green")

    if not yes:
        click.confirm(click.style("Do you want to continue?", bold=True), abort=True)

    run(
        [
            "docker",
            "compose",
            "--project-directory",
            ctx.obj["manifests_path"] / service,
            "pull",
        ],
        check=True,
    )

    if service == "discovery-provider" and (seed or auto_seed):
        # make sure services that use postgres are not running
        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / "discovery-provider",
                "down",
                "backend",
                "indexer",
                "comms",
                "notifications",
                "es-indexer",
            ],
        )

        click.secho("Seeding the discovery provider", fg="yellow")
        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / "discovery-provider",
                "--profile",
                "discovery",
                "run",
                "seed",
                "bash",
                "/usr/share/seed.sh",
                get_network(ctx, "discovery-provider"),
                "true" if auto_seed else "false",
            ],
        )

    plugins = get_registered_plugins(ctx, service)

    run(
        [
            "docker",
            "compose",
            "--project-directory",
            ctx.obj["manifests_path"] / service,
            *plugins,
            "up",
            "--remove-orphans",
            "-d",
        ],
        check=True,
    )

    prune()


@cli.command()
@click.argument("service", type=service_type)
@click.argument("containers", type=container_type, nargs=-1)
@click.pass_context
def logs(ctx, service, containers):
    """Get logs for a service/container"""
    run(
        [
            "docker",
            "compose",
            "--project-directory",
            ctx.obj["manifests_path"] / service,
            "logs",
            "-f",
            *containers,
        ],
    )


@cli.command()
@click.argument("service", type=service_type, required=False)
@click.argument("containers", type=container_type, nargs=-1)
@click.pass_context
def restart(ctx, service, containers):
    """Restart a service/container"""
    set_automatic_env(ctx)

    lock(ctx, "docker")

    services = [service]
    if service is None:
        services = SERVICES

    for service in services:
        server = "backend"
        if service == "creator-node":
            server = "audiusd"
        proc = run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                "ps",
                "-q",
                server,
            ],
            capture_output=True,
        )

        plugins = get_registered_plugins(ctx, service)

        if proc.stdout:
            run(
                [
                    "docker",
                    "compose",
                    "--project-directory",
                    ctx.obj["manifests_path"] / service,
                    *plugins,
                    "up",
                    "--build",
                    "--force-recreate",
                    "--remove-orphans",
                    "-d",
                    *containers,
                ],
            )

    prune()


@cli.command()
@click.argument("service", type=service_type, required=False)
@click.argument("containers", type=container_type, nargs=-1)
@click.pass_context
def down(ctx, service, containers):
    """Stops a service/container"""
    lock(ctx, "docker")

    services = [service]
    if service is None:
        services = SERVICES

    for service in services:
        if not containers and service == "discovery-provider":
            # ensure removal of the chain container
            run(
                [
                    "docker",
                    "compose",
                    "--project-directory",
                    ctx.obj["manifests_path"] / service,
                    "--profile",
                    "chain",
                    "down",
                ],
            )

        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                "down",
                *containers,
            ],
        )


@cli.command()
@click.option("--unset", is_flag=True)
@click.option("--required", is_flag=True)
@click.argument("service", type=service_type)
@click.argument("key", required=False)
@click.argument("value", required=False)
@click.pass_context
def set_config(ctx, service, unset, required, key, value):
    """Set a config value"""
    if required and (key or value):
        click.secho("--required cannot be used when key or value is set", fg="red")
        sys.exit(1)

    env = ctx.obj["manifests_path"] / service / f"{get_network(ctx, service)}.env"
    env_data = dotenv.dotenv_values(env)

    override_env = get_override_path(ctx, service)
    override_env_data = dotenv.dotenv_values(override_env)

    if required:
        logging.info("audius-cli set-config required")
        for key, value in env_data.items():
            if not unset and value == "":
                value = click.prompt(
                    click.style(key, bold=True),
                    override_env_data.get(key, env_data.get(key)),
                )
                logging.info(f"> audius-cli set-config key={key!r} value={value!r}")
                dotenv.set_key(override_env, key, value)
            if unset and value == "":
                logging.info(f"> audius-cli set-config unset key={key!r}")
                dotenv.unset_key(override_env, key)
    else:
        if key is None:
            key = click.prompt(click.style("Key", bold=True))
        if not unset and value is None:
            value = click.prompt(click.style("Value", bold=True))

        if unset:
            logging.info(f"audius-cli set-config unset key={key!r}")
            dotenv.unset_key(override_env, key)
        else:
            logging.info(f"audius-cli set-config key={key!r} value={value!r}")
            dotenv.set_key(override_env, key, value)


@cli.command()
@click.argument("service", type=service_type)
@click.pass_context
def get_config(ctx, service):
    server = "backend"
    if service == "creator-node":
        server = "audiusd"
    proc = run(
        [
            "docker",
            "compose",
            "--project-directory",
            ctx.obj["manifests_path"] / service,
            "ps",
            "-q",
            server,
        ],
        capture_output=True,
    )

    # get from inside container if running otherwise fallback override.env
    if proc.stdout:
        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                "exec",
                server,
                "env",
            ]
        )
    else:
        env = ctx.obj["manifests_path"] / service / f"{get_network(ctx, service)}.env"
        env_data = dotenv.dotenv_values(env)
        override_env = get_override_path(ctx, service)
        override_env_data = dotenv.dotenv_values(override_env)
        current_env_data = {**env_data, **override_env_data}
        for key, value in current_env_data.items():
            print(f"{key}={value}")


@cli.command()
@click.option("--unset", is_flag=True)
@click.option("-y", "--yes", is_flag=True)
@click.argument("tag", required=False)
@click.pass_context
def set_tag(ctx, unset, yes, tag):
    """Set the commit tag"""
    if not unset and tag is None:
        tag = click.prompt(click.style("Tag", bold=True))

    if not unset:
        try:
            for commit in json.load(
                urllib.request.urlopen(
                    f"https://api.github.com/repos/AudiusProject/audius-protocol/commits?sha={tag}&per_page=10"
                )
            ):
                short_message = commit["commit"]["message"].split("\n")[0]
                click.echo(
                    f"{click.style(commit['sha'][:8], bold=True, fg='yellow')}: "
                    f"[{click.style(commit['commit']['author']['name'], bold=True, fg='blue')}] "
                    f"{short_message}"
                )
        except (ConnectionError, urllib.error.HTTPError, json.JSONDecodeError):
            click.secho("Failed to get commit messages", fg="red")

    logging.info(f"audius-cli set-tag tag={tag!r}")
    for service in SERVICES:
        key = "TAG"
        env_file = ctx.obj["manifests_path"] / service / ".env"
        if unset:
            logging.info(f"> audius-cli set-tag unset service={service!r}")
            dotenv.unset_key(env_file, key)
        else:
            logging.info(f"> audius-cli set-tag service={service!r} tag={tag!r}")
            dotenv.set_key(env_file, key, tag)


@cli.command()
@click.option("--unset", is_flag=True)
@click.argument("service", required=True)
@click.argument("path", type=click.Path(), required=False)
@click.pass_context
def set_override_path(ctx, unset, service, path):
    """Specify alternate location for override.env"""
    if not unset:
        if not service:
            service = click.prompt(click.style("Service", bold=True))
        if not path:
            path = click.prompt(click.style("Path", bold=True))

    env_file = ctx.obj["manifests_path"] / service / ".env"

    if unset:
        logging.info(
            f"audius-cli set-override-path unset service={service!r} path={path!r}"
        )
        dotenv.unset_key(env_file, "OVERRIDE_PATH")
    else:
        logging.info(f"audius-cli set-override-path service={service!r} path={path!r}")
        dotenv.set_key(env_file, "OVERRIDE_PATH", path)


@cli.command()
@click.option("--unset", is_flag=True)
@click.argument("network", type=network_type, required=False)
@click.pass_context
def set_network(ctx, unset, network):
    """Set the deployment network"""
    if not unset and network is None:
        network = click.prompt(click.style("Network", bold=True))

    logging.info(f"audius-cli set-network network={network!r}")
    for service in SERVICES:
        env_file = ctx.obj["manifests_path"] / service / ".env"

        if unset:
            logging.info(f"> audius-cli set-network unset service={service!r}")
            dotenv.unset_key(env_file, "NETWORK")
        else:
            logging.info(
                f"> audius-cli set-network service={service!r} network={network!r}"
            )
            dotenv.set_key(env_file, "NETWORK", network)


@cli.command()
@click.argument("branch", required=False)
@click.pass_context
def pull_reset(ctx, branch):
    """Pull latest updates from remote and hard reset"""
    try:
        subprocess.run(["git", "fetch"], check=True, cwd=ctx.obj["manifests_path"])

        if branch:
            subprocess.run(
                ["git", "checkout", branch], check=True, cwd=ctx.obj["manifests_path"]
            )

        # Get the name of the current branch
        current_branch = subprocess.check_output(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
            cwd=ctx.obj["manifests_path"],
            text=True,
        ).strip()

        # Run git reset with the current branch
        subprocess.run(
            ["git", "reset", "--hard", f"origin/{current_branch}"],
            check=True,
            cwd=ctx.obj["manifests_path"],
        )

    except subprocess.CalledProcessError as e:
        click.secho(f"Could not pull and hard reset: {e}", fg="red")
        sys.exit(1)


@cli.command()
@click.argument("branch", required=False)
@click.pass_context
def upgrade(ctx, branch):
    """Pulls from latest source and re-launches all running services"""
    log_file = ctx.obj["manifests_path"] / "auto-upgrade.log"
    try:
        # address bug where `auto-upgrade.log: Is a directory`
        if log_file.is_dir():
            shutil.rmtree(log_file)
        open(log_file, "w").close()
    except Exception as e:
        print(f"Error while clearing the log file: {e}")

    try:
        ctx.forward(pull_reset)
        set_automatic_env(ctx)

        lock(ctx, "docker")

        identity_override_env = get_override_path(ctx, "identity-service")
        creator_override_env = get_override_path(ctx, "creator-node")
        disc_override_env = get_override_path(ctx, "discovery-provider")
        
        # determine service type based on override.env location and env vars in it
        service = ""
        if creator_override_env.exists() and dotenv.get_key(creator_override_env, "creatorNodeEndpoint"):
            service = "creator-node"
        elif identity_override_env.exists() and dotenv.get_key(identity_override_env, "ethOwnerWallet"):
            service = "identity-service"
        elif disc_override_env.exists() and dotenv.get_key(disc_override_env, "audius_delegate_owner_wallet"):
            service = "discovery-provider"
        else:
            click.secho("Unable to determine service type. Ensure your override.env is configured", fg="yellow")
            sys.exit(1)

        if service == "discovery-provider" and not is_in_discprov_whitelist(ctx):
            # Stop and remove all containers except vector and audiusd
            containers = run(
                ["docker", "ps", "--format", "{{.Names}}"],
                capture_output=True,
                text=True
            ).stdout.strip().split('\n')
            containers = [c for c in containers if c and c not in ('vector', 'audiusd')]
            if containers:  # Only run if containers were found
                run(["docker", "rm", "-f"] + containers)
        plugins = get_registered_plugins(ctx, service)

        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                *plugins,
                "pull",
            ],
        )

        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                *plugins,
                "up",
                "--remove-orphans",
                "-d",
            ],
        )

        run(
            [
                "docker",
                "compose",
                "--project-directory",
                ctx.obj["manifests_path"] / service,
                *plugins,
                "up",
                "--force-recreate",
                "vector",
                "--remove-orphans",
                "-d",
            ],
        )

    finally:
        lock(ctx, "docker")
        prune()


@cli.command()
@click.option("--remove", is_flag=True)
# random so nodes in the network stagger upgrades
@click.argument("cron-expression", default=f"{random.randint(0, 59)} * * * *")
@click.pass_context
def auto_upgrade(ctx, remove, cron_expression):
    """Setup auto upgrade with a cron job"""
    with crontab.CronTab(user=True) as cron:
        log_file = ctx.obj["manifests_path"] / "auto-upgrade.log"
        
        for job in cron.find_comment("audius-cli auto-upgrade"):
            print(
                f'Removed auto-upgrade cron. If you want to re-enable, run: audius-cli auto-upgrade "{job.slices}"'
            )
            cron.remove(job)
            # remove any lingering auto-upgrade.log directory that should be a file
            if log_file.is_dir():
                shutil.rmtree(log_file)

        if not remove:
            try:
                open(log_file, "x").close()
            except FileExistsError:
                pass
            except:
                logging.warn(f"Unable to create `{log_file}` file")

            job = cron.new(
                (
                    f"/usr/local/bin/audius-cli upgrade >> {log_file} 2>&1;"
                ),
                "audius-cli auto-upgrade",
            )

            job.setall(cron_expression)


@cli.command()
@click.pass_context
@click.argument("service", type=service_type)
@click.argument("name", required=True)
def register_plugin(ctx, service, name):
    print(f"registering plugin {name} on {service}")
    override_env = get_override_path(ctx, service)
    plugins = dotenv.get_key(override_env, REGISTERED_PLUGINS)
    if plugins is not None:
        plugins = plugins.split(",")
    else:
        plugins = []
    plugins.append(name)
    plugins = set(filter(lambda x: x != "", plugins))
    plugins = ",".join(plugins)
    dotenv.set_key(override_env, REGISTERED_PLUGINS, plugins)
    return


@cli.command()
@click.pass_context
@click.argument("service", type=service_type)
@click.argument("name", required=True)
def deregister_plugin(ctx, service, name):
    print(f"deregistering plugin {name} on {service}")
    override_env = get_override_path(ctx, service)
    plugins = dotenv.get_key(override_env, REGISTERED_PLUGINS)
    if plugins is not None:
        plugins = plugins.split(",")
    else:
        return
    if name in plugins:
        plugins.remove(name)
    plugins = set(filter(lambda x: x != "", plugins))
    plugins = ",".join(plugins)
    dotenv.set_key(override_env, REGISTERED_PLUGINS, plugins)
    return


@cli.command()
@click.pass_context
@click.argument("service", type=service_type)
def print_plugins(ctx, service):
    override_env = get_override_path(ctx, service)
    plugins = dotenv.get_key(override_env, REGISTERED_PLUGINS)
    print(f"registered plugins on {service}: {plugins}")
    return


def get_registered_plugins(ctx, service):
    """
    Formats registered plugins into a "--profile {plugin}" format
    so it can be passed into an up command.
    """
    override_env = get_override_path(ctx, service)
    plugins_env = dotenv.get_key(override_env, REGISTERED_PLUGINS)

    registered_plugins = []
    if plugins_env is not None:
        registered_plugins = plugins_env.split(",")

    if not registered_plugins:
        return []

    profiles = []
    for plugin in registered_plugins:
        profiles.extend(["--profile", plugin])

    if service == "discovery-provider" and is_in_discprov_whitelist(ctx):
        profiles.extend(["--profile", "discovery"])
    else:
        profiles.extend(["--profile", "audiusd"])

    return profiles


def is_in_discprov_whitelist(ctx):
    override_env = get_override_path(ctx, "discovery-provider")
    url_env = dotenv.get_key(override_env, "audius_discprov_url")
    for suffix in USE_LEGACY_DISCOVERY_PROVIDER_PROFILE:
        if url_env.endswith(suffix):
            return True
    return False

if __name__ == "__main__":
    cli(obj={})