audius-compose

#!/usr/bin/env python3

import getpass
import json
import os
import pathlib
import subprocess
import sys
import urllib.request

import click
import dotenv
import eth_account
import solders.keypair

LOCALHOSTS = """
127.0.0.1       ingress \
audius-protocol-comms-discovery-1 audius-protocol-comms-discovery-2 audius-protocol-comms-discovery-3 audius-protocol-comms-discovery-4 audius-protocol-comms-discovery-5 \
audius-protocol-discovery-provider-1 audius-protocol-discovery-provider-2 audius-protocol-discovery-provider-3 \
audius-protocol-creator-node-1 audius-protocol-creator-node-2 audius-protocol-creator-node-3 \
audius-protocol-identity-service-1 \
audius-protocol-solana-test-validator-1 \
audius-protocol-eth-ganache-1 \
audius-protocol-poa-ganache-1 \
audius-protocol-pedalboard \
audius-protocol-anti-abuse-oracle-1 \
audius-protocol-discovery-provider-redis-1 \
"""


class DefaultCommandGroup(click.Group):
    """allow a default command for a group"""

    def command(self, *args, **kwargs):
        default_command = kwargs.pop("default_command", False)
        if default_command and not args:
            kwargs["name"] = kwargs.get("name", "<>")
        decorator = super(DefaultCommandGroup, self).command(*args, **kwargs)

        if default_command:

            def new_decorator(f):
                cmd = decorator(f)
                self.default_command = cmd.name
                return cmd

            return new_decorator

        return decorator

    def resolve_command(self, ctx, args):
        try:
            # test if the command parses
            return super(DefaultCommandGroup, self).resolve_command(ctx, args)
        except click.UsageError:
            # command did not parse, assume it is the default command
            args.insert(0, self.default_command)
            return super(DefaultCommandGroup, self).resolve_command(ctx, args)


def get_git_commit(protocol_dir):
    return (
        subprocess.run(
            ["git", f"--git-dir={protocol_dir / '.git'}", "rev-parse", "HEAD"],
            capture_output=True,
        )
        .stdout.decode()
        .strip()
    )


# TODO: Maybe clean up later to handle comms config (currently using explicit env vars in docker-compose.comms.yml)
def generate_env(
    protocol_dir, discovery_provider_replicas, elasticsearch_replicas, prod=False
):
    config = {}
    if (protocol_dir / "dev-tools/config.json").exists():
        config = json.load((protocol_dir / "dev-tools/config.json").open())

    env = {}

    env["PROJECT_ROOT"] = str(protocol_dir)

    env["GIT_COMMIT"] = get_git_commit(protocol_dir)

    env["DISCOVERY_PROVIDER_REPLICAS"] = str(discovery_provider_replicas)

    env["CONTENT_NODE_VERSION"] = json.loads(
        (protocol_dir / "pkg/mediorum/.version.json").read_text(),
    )["version"]

    env["DISCOVERY_NODE_VERSION"] = json.loads(
        (protocol_dir / "packages/discovery-provider/.version.json").read_text(),
    )["version"]

    env["ELASTICSEARCH_REPLICAS"] = str(elasticsearch_replicas)
    if elasticsearch_replicas:
        env["ELASTICSEARCH_CONDITION"] = "service_healthy"
    else:  # exists to prevent discovery provider from waiting for non-existent elasticsearch instances
        env["ELASTICSEARCH_CONDITION"] = "service_started"

    for name, secret_key in config.get("solana-accounts", {}).items():
        keypair = solders.keypair.Keypair.from_bytes(bytes(secret_key))
        env[f"{name}_SECRET_KEY"] = json.dumps(list(bytes(keypair)))
        env[f"{name}_PUBLIC_KEY"] = str(keypair.pubkey())

    eth_relayer_wallets = []
    poa_relayer_wallets = []

    for name, private_key in config.get("eth-accounts", {}).items():
        account = eth_account.Account.from_key(private_key)

        # slice off 0x for hex private keys
        env[f"{name}_PRIVATE_KEY"] = account.key.hex().replace("0x", "")
        env[f"{name}_ADDRESS"] = account.address

        if name.startswith("ETH_RELAYER_WALLET"):
            eth_relayer_wallets.append(
                {
                    "publicKey": account.address,
                    "privateKey": account.key.hex().replace("0x", ""),
                }
            )

        if name.startswith("POA_RELAYER_WALLET"):
            poa_relayer_wallets.append(
                {
                    "publicKey": account.address,
                    "privateKey": account.key.hex().replace("0x", ""),
                }
            )

    env["ETH_RELAYER_WALLETS"] = json.dumps(eth_relayer_wallets)
    env["POA_RELAYER_WALLETS"] = json.dumps(poa_relayer_wallets)

    bootstrap_sp_ids = []
    bootstrap_sp_owner_wallets = []
    for replica in range(6):
        if f"CN{replica + 1}_SP_OWNER_ADDRESS" in env:
            bootstrap_sp_ids.append(str(replica + 1))
            bootstrap_sp_owner_wallets.append(env[f"CN{replica + 1}_SP_OWNER_ADDRESS"])
    env["BOOTSTRAP_SP_IDS"] = ",".join(bootstrap_sp_ids)
    env["BOOTSTRAP_SP_OWNER_WALLETS"] = ",".join(bootstrap_sp_owner_wallets)
    env["BOOTSTRAP_SP_DELEGATE_WALLETS"] = ",".join(bootstrap_sp_owner_wallets)

    aao_wallets = []
    for key, value in env.items():
        if key.startswith("AAO_WALLET_") and key.endswith("_ADDRESS"):
            aao_wallets.append(value)
    env["AAO_WALLET_ADDRESSES"] = ",".join(aao_wallets)

    if prod:
        env["DOCKERCOMPOSE_ENV_TYPE"] = "prod"
    else:
        env["DOCKERCOMPOSE_ENV_TYPE"] = "dev"

    for key, value in config.get("extra-env", {}).items():
        env[key] = value

    # generate config.env used by startup scripts
    env_file = protocol_dir / "dev-tools/compose/.env"
    env_file.touch()
    env_file.write_text(
        "# DO NOT EDIT THIS FILE\n# Instead edit dev-tools/config.json and dev-tools/audius-compose.py generate_env()\n"
    )
    for key, value in env.items():
        dotenv.set_key(env_file, key, value)


@click.group()
@click.option(
    "--protocol-dir",
    type=click.Path(
        exists=True, file_okay=False, resolve_path=True, path_type=pathlib.Path
    ),
)
@click.pass_context
def cli(ctx, protocol_dir):
    if protocol_dir is None:
        protocol_dir = pathlib.Path(
            os.path.dirname(os.path.realpath(__file__))
        ).resolve()
        while (
            protocol_dir.name != ""
            and not (protocol_dir / "dev-tools/compose/docker-compose.yml").exists()
        ):
            protocol_dir = protocol_dir.parent
        if protocol_dir.name == "":
            raise click.ClickException("Unable to find protocol dir")

    ctx.obj = protocol_dir


@cli.command()
@click.option("-d", "--discovery-provider-replicas", default=1, type=int)
@click.option("-e", "--elasticsearch-replicas", default=0, type=int)
@click.option("-a", "--args", type=str, multiple=True)
@click.option(
    "--prod",
    is_flag=True,
    help="Use production containers which have all dependencies built in",
)
@click.argument("services", nargs=-1)
@click.pass_obj
def build(
    protocol_dir,
    discovery_provider_replicas,
    elasticsearch_replicas,
    prod,
    args,
    services,
):
    generate_env(
        protocol_dir, discovery_provider_replicas, elasticsearch_replicas, prod
    )

    proc = subprocess.run(
        [
            "docker",
            "compose",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.yml'}",
            f"--project-name={protocol_dir.name}",
            "--profile=circle_ci_push",
            "--profile=tests",
            "--profile=identity",
            "--profile=discovery",
            "--profile=comms",
            "--profile=storage",
            "--profile=notifications",
            "--profile=elasticsearch",
            "--profile=mediorum",
            "--profile=chain",
            "--profile=poa",
            "--profile=eth",
            "--profile=libs",
            "--profile=solana",
            "--profile=block-explorer",
            "build",
            *args,
            *services,
        ],
    )

    if proc.returncode != 0:
        raise click.ClickException("Failed to build images")


@cli.command()
@click.argument("services", nargs=-1)
@click.option(
    "--prod",
    is_flag=True,
    help="Use production containers which have all dependencies built in",
)
@click.pass_context
def push(ctx, services, prod):
    protocol_dir = ctx.obj
    services = services or ["mediorum", "discovery-provider", "identity-service"]

    git_commit = get_git_commit(protocol_dir)

    if not git_commit:
        raise click.ClickException("Failed to get git commit")

    # this is building w/o cache every time
    ctx.invoke(build, services=services, prod=prod)

    try:
        for service in services:
            base_tag = f"{protocol_dir.name}-{service}"

            subprocess.run(
                ["docker", "tag", base_tag, f"audius/{service}:{git_commit}"],
                check=True,
            )
            subprocess.run(
                ["docker", "push", f"audius/{service}:{git_commit}"], check=True
            )
    except subprocess.CalledProcessError:
        raise click.ClickException("Failed to tag images")

    try:
        for service in services:
            subprocess.run(
                ["docker", "push", f"audius/{service}:{git_commit}"],
                check=True,
            )
    except subprocess.CalledProcessError:
        raise click.ClickException("Failed to push images")


@cli.command()
@click.pass_context
def pull(ctx):
    protocol_dir = ctx.obj
    subprocess.run(
        [
            "docker",
            "compose",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.yml'}",
            f"--project-name={protocol_dir.name}",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "--profile=solana",
            "pull",
            "--ignore-buildable",
        ],
    )


@cli.command()
@click.argument("service")
@click.argument("name")
@click.argument("value")
@click.pass_obj
def set_env(protocol_dir, service, name, value):
    env_file = protocol_dir / f"dev-tools/startup/{service}.env"
    dotenv.set_key(env_file, name, value)


@cli.command()
@click.argument("service")
@click.argument("environment")
@click.pass_obj
def load_env(protocol_dir, service, environment):
    if not click.confirm(f"All existing env for {service} will be replaced continue?"):
        return

    replica = 1
    parts = service.split("-")
    if parts[-1].isdigit():
        service, replica = "-".join(parts[:-1]), int(parts[-1])

    env_file = protocol_dir / f"dev-tools/startup/{service}-{replica}.env"

    with urllib.request.urlopen(
        f"https://raw.githubusercontent.com/AudiusProject/audius-docker-compose/main/{service}/{environment}.env"
    ) as resp:
        env_file.write_bytes(resp.read())


@cli.command()
@click.option("-d", "--discovery-provider-replicas", default=1, type=int)
@click.option("-e", "--elasticsearch-replicas", default=1, type=int)
@click.option("-b", "--block-explorer", is_flag=True)
@click.option("-o", "--anti-abuse-oracle", is_flag=True)
@click.option(
    "-n",
    "--notifs",
    is_flag=True,
    help="Bring up discovery provider notifications container",
)
@click.option(
    "-p",
    "--pedalboard",
    is_flag=True,
    help="Bring up the pedalboard containers",
)
@click.option("-a", "--args", type=str, multiple=True)
@click.option(
    "--prod",
    is_flag=True,
    help="Use production containers which have all dependencies built in",
)
@click.argument("services", nargs=-1)
@click.pass_obj
@click.pass_context
def up(
    ctx,
    protocol_dir,
    discovery_provider_replicas,
    elasticsearch_replicas,
    block_explorer,
    anti_abuse_oracle,
    notifs,
    pedalboard,
    args,
    prod,
    services,
):
    generate_env(
        protocol_dir, discovery_provider_replicas, elasticsearch_replicas, prod
    )

    AAO_DIR = pathlib.Path(os.getenv("AAO_DIR", protocol_dir / "../anti-abuse-oracle"))
    os.environ["MEDIORUM_ENV_TYPE"] = (
        "dev"  # Mediorum's Dockerfile.prod doesn't support multiple mediorum instances, so always use dev
    )
    profiles = (
        *(["--profile=block-explorer"] if block_explorer else []),
        *(["--profile=elasticsearch"] if elasticsearch_replicas else []),
        *(["--profile=notifications"] if notifs else []),
        *(["--profile=pedalboard"] if pedalboard else []),
        "--profile=storage",
        "--profile=libs",
        "--profile=solana",
        "--profile=identity",
        "--profile=discovery",
        "--profile=comms",
        "--profile=chain",
        *(
            ["--file=" + str(AAO_DIR / "audius-compose.yml")]
            if anti_abuse_oracle
            else []
        ),
    )

    ctx.invoke(pull)
    sys.exit(
        subprocess.run(
            [
                "docker",
                "compose",
                f"--project-directory={protocol_dir / 'dev-tools/compose'}",
                f"--project-name={protocol_dir.name}",
                f"--file={protocol_dir / 'dev-tools/compose/docker-compose.yml'}",
                *profiles,
                "up",
                "--build",
                "--detach",
                "--wait",
                *args,
                *services,
            ],
        ).returncode
    )


@cli.command()
@click.pass_obj
def down(protocol_dir):
    subprocess.run(
        [
            "docker",
            "compose",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            f"--project-name={protocol_dir.name}",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.yml'}",
            "--profile=*",
            "down",
            "--remove-orphans",
            "--volumes",
        ],
    )


@cli.command()
@click.option("-n", "--nopass", is_flag=True, help="Use sudo without password (for CI)")
def connect(nopass):
    hosts_file = pathlib.Path("/etc/hosts")
    if LOCALHOSTS in hosts_file.read_text():
        print("/etc/hosts already up to date")
    else:
        print("Adding development hosts to /etc/hosts.")
        password = "" if nopass else getpass.getpass("Enter your sudo password: ")
        subprocess.run(
            f"echo '{password}' | sudo -S sh -c 'echo \"{LOCALHOSTS}\" >> /etc/hosts'",
            shell=True,
        )
        print("Updated /etc/hosts")


@cli.group(cls=DefaultCommandGroup)
def test():
    pass


@test.command(name="build")
@click.pass_obj
def test_build(protocol_dir):
    generate_env(protocol_dir, 0, 0)

    subprocess.run(
        [
            "docker",
            "compose",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.test.yml'}",
            f"--project-name={protocol_dir.name}-test",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "build",
        ],
    )

    subprocess.run(
        [
            "docker",
            "compose",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.test.yml'}",
            f"--project-name={protocol_dir.name}-test",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "pull",
        ],
    )


@test.command(name="run", default_command=True)
@click.argument("service")
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
@click.pass_obj
def test_run(protocol_dir, service, args):
    """
    Runs tests for a given service.
    """
    generate_env(protocol_dir, 0, 0)
    env = os.environ.copy()

    result = subprocess.run(
        [
            "docker",
            "compose",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.test.yml'}",
            f"--project-name={protocol_dir.name}-test",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "run",
            "--rm",
            "--build",
            f"test-{service}",
            "test",
            *args,
        ],
        env=env,
    )

    sys.exit(result.returncode)


@test.command(name="up")
@click.argument("service")
@click.pass_obj
def test_up(protocol_dir, service):
    """
    Brings up a testing container for a given service.
    """
    generate_env(protocol_dir, 0, 0)

    sys.exit(
        subprocess.run(
            [
                "docker",
                "compose",
                f"--file={protocol_dir / 'dev-tools/compose/docker-compose.test.yml'}",
                f"--project-name={protocol_dir.name}-test",
                f"--project-directory={protocol_dir / 'dev-tools/compose'}",
                "up",
                "-d",
                "--build",
                f"test-{service}",
            ]
        ).returncode
    )


@test.command(name="down")
@click.pass_obj
def test_down(protocol_dir):
    generate_env(protocol_dir, 0, 0)

    subprocess.run(
        [
            "docker",
            "compose",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.test.yml'}",
            f"--project-name={protocol_dir.name}-test",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "--profile=*",
            "down",
            "-v",
        ],
    )


@cli.command()
@click.argument("service")
@click.argument("command")
@click.argument("args", nargs=-1)
@click.pass_obj
def exec(protocol_dir, service, command, args):
    index = 1
    split = service.rsplit("-", 1)
    if split[-1].isdigit():
        service, index = split

    subprocess.run(
        [
            "docker",
            "compose",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            "exec",
            "--index",
            str(index),
            service,
            "sh",
            "-c",
            # Ugly command since this needs run on all containers (i.e. posix compliant)
            r'''eval $(od -An -to1 -v /proc/1/environ | tr -d '\n' | sed 's/ 000/\n/g; s/ /\\/g; s/[^\n]\+/export "`printf "&"`"/g'); exec "$@"''',
            # r'''eval $(xargs -0 -n1 sh -c 'echo $0 | sed "'"s/'/'\\\\\''/g; s/.*/export '&'/"'"' </proc/1/environ); exec "$@"''',
            "-",
            command,
            *args,
        ],
    )


@cli.command()
@click.pass_obj
def ps(protocol_dir):
    proc = subprocess.run(
        [
            "docker",
            "ps",
            "--format",
            "{{json .}}",
        ],
        capture_output=True,
        text=True,
    )

    if proc.returncode != 0:
        raise click.ClickException(proc.stderr.decode())

    output = proc.stdout.strip()
    services = [json.loads(line) for line in output.splitlines()]
    services.sort(key=lambda x: x["Names"])

    print(
        "CONTAINER ID".ljust(13),
        "NAME".ljust(35),
        "STATUS".ljust(10),
        "PORTS",
    )

    for service in services:
        name = service["Names"]
        status = service["Status"] or service["State"]

        ports = service["Ports"]

        if service["Names"] in [
            "discovery-provider",
            "discovery-provider-elasticsearch",
        ]:
            replica = int(service["Name"].split("-")[-1])
        if service["Names"] == "discovery-provider":
            name = f"{service['Names']}-{replica}"
            ports[5000 + replica - 1] = 5000
        if service["Names"] == "discovery-provider-elasticsearch":
            name = f"{service['Names']}-{replica}"

        name = name.removeprefix("audius-protocol-")

        print(
            service["ID"][:12].ljust(13),
            name.ljust(35),
            status.ljust(10),
            ports,
        )

    print("** most containers prefixed with 'audius-protocol-' **")


@cli.command()
@click.option("-y", "--yes", is_flag=True)
@click.pass_obj
def prune(protocol_dir, yes):
    click.secho(
        "WARNING! This will remove all dangling images and build cache not used within the last 72h",
        fg="yellow",
    )

    if yes or click.confirm("Are you sure you want to continue?"):
        subprocess.run(["docker", "image", "prune", "-f"])
        subprocess.run(["docker", "builder", "prune", "-f", "--keep-storage=30G"])


@cli.command()
@click.argument("service")
@click.pass_obj
def logs(protocol_dir, service):
    subprocess.run(
        [
            "docker",
            "compose",
            f"--project-directory={protocol_dir / 'dev-tools/compose'}",
            f"--file={protocol_dir / 'dev-tools/compose/docker-compose.yml'}",
            f"--project-name={protocol_dir.name}",
            "logs",
            service,
        ],
    )


if __name__ == "__main__":
    cli()