mob

#!/usr/bin/env python3
# Copyright (c) 2018-2022 The MobileCoin Foundation
"""
A helper tool for getting the dockerized build environment

`mob` tool helps you get a prompt in the same Docker environment used in CI,
to make building easier.

operation
---------

The basic operation is like this:

0. `./mob prompt` is invoked

1. Read .mobconf to find the remote source for the image.

1. Do an appropriate form of `docker run -it bash` in this image, mounting the
root of the repository to `/tmp/mobilenode` and setting that as the working directory.

There are some flags and options for modifying this process, e.g. `--dry-run` just
shows you the shell commands without executing them.

`mob` tool attempts to mount the `/dev/isgx` device correctly if it is available
and you selected `--hw` mode, so that you can run tests in hardware mode.

usage notes (ssh)
-----------------

The `--ssh-dir` and `--ssh-agent` options can be used if the build requires
access to private repos on github (mobilecoin repo sources will be mounted so
ssh stuff is not needed for that, this is for any private dependencies of the mobilecoin
repo.) If provided, these options will try to get your credentials from the host
environment into the container so that cargo can pull.

mobconf
-------

`mob` supports configuration via the `.mobconf` file. This allows it to build multiple
different projects that exist in the same repository.

`mob` attempts to find a `.mobconf` file by starting in `pwd` and searching up,
it is an error if it can't find one.

.mobconf sections:

[image]
url = The image URL to pass to docker run
tag = Image tag/version to use
"""

import argparse
import configparser
import getpass
import grp
import os
import pathlib
import platform
import subprocess
import sys

parser = argparse.ArgumentParser(prog="mob", description="Perform an action or get a prompt in docker build environment")
parser.add_argument("action", choices=["prompt", "run"], help="""
                    (prompt) Run bash in the build environment
                    (run) Run the given command (passed to `docker run`)
                    """)
parser.add_argument("--name", nargs='?', default=None, help="The name for the container, run with prompt")
parser.add_argument("--hw", action="store_true", help="Set SGX_MODE=HW. Default is SGX_MODE=SW")
parser.add_argument("--ias-prod", action="store_true", help="Set IAS_MODE=PROD. Default is IAS_MODE=DEV. This affects which IAS endpoints we use.")
parser.add_argument("--dry-run", action="store_true", help="Don't run docker, show how we would invoke docker.")
parser.add_argument("--no-pull", action='store_true', help='Skip the `docker image pull` step.')
parser.add_argument("--tag", default=None, type=str, help="Use given tag for image rather than the one in .mobconf")
parser.add_argument("--verbose", action="store_true", help="Show the commands on stdout. True by default when noninteractive, implied by dry-run")
parser.add_argument("--ssh-dir", action='store_true', help='Mount $HOME/.ssh directory into /root/.ssh for ssh key access')
parser.add_argument("--ssh-agent", action='store_true', help='Use ssh-agent on host machine for ssh authentication. Also controlled by SSH_AUTH_SOCK env variable')
parser.add_argument("--expose", nargs='+', default=None, help="Any additional ports to expose")
parser.add_argument("--publish", nargs='+', default=None, help="Any additional ports to publish, e.g. if running wallet locally.")
parser.add_argument("--cmd", nargs='+', default=None, help="Run this command inside the bash shell instead of an interactive prompt")
parser.add_argument("--run-as-root", action='store_true', help="run prompt as root instead of local user")

args = parser.parse_args()

###
# Implement forced settings
###
if args.dry_run or not sys.stdout.isatty():
    args.verbose = True

###
# Implement verbose, dry_run settings
###

def eprint(*argv, **kwargs):
    """
    When python is invoked from docker in CI, we won't see anything because of
    buffered output unless we flush. It's hard to ensure PYTHONUNBUFFERED=0 or -u
    is used consistently.
    """
    print(*argv, file=sys.stderr, **kwargs)
    sys.stderr.flush()

def vprint(*argv, **kwargs):
    """ vprint is eprint that only happens in verbose mode """
    if args.verbose:
        eprint(*argv, **kwargs)

def vprint_command(cmd):
    """ Print a command, whether it is in shell style or list style. """
    if isinstance(cmd, list):
        cmd = ' '.join(cmd)
    vprint(f'$ {cmd}')

# Run a command, unless we are in dry_run mode and should not do that
# Print the command if in verbose mode
def maybe_run(cmd, **kwargs):
    vprint_command(cmd)
    if not args.dry_run:
        return subprocess.check_call(cmd, **kwargs)

# Change directory.
# Print the command if in verbose mode
def verbose_chdir(path):
    vprint_command(['cd', path])
    os.chdir(path)

# Check if we have a git commit and compute outside the container
# Sometimes git cannot be used in the container, if building in public dir,
# or if you used git worktree to make a new worktree.
def get_git_commit():
    if "GIT_COMMIT" in os.environ:
        return os.environ["GIT_COMMIT"]
    else:
        try:
            cmd = ["git", "describe", "--always", "--dirty=-modified"]
            vprint_command(cmd)
            return subprocess.check_output(cmd)[:-1].decode()
        except subprocess.CalledProcessError:
            eprint("Couldn't get git revision")
            return None

##
# Environment checks
##

# Check if docker is available and bail out early with a message if appropriate
if not args.dry_run:
    maybe_run("command -v docker > /dev/null 2>&1", shell=True)

# Check for any SSH handling
if "SSH_AUTH_SOCK" in os.environ:
    if platform.system().lower() != "linux":
        vprint("SSH Auth Socket found, but not running on linux")
        args.ssh_dir = True
    else:
        vprint("Mapping SSH_AUTH_SOCKET into container")
        args.ssh_agent = True

# Find work directory and change directory there
# This is based on searching for nearest .mobconf file, moving upwards from CWD
top_level = os.getcwd()
while not os.path.exists(os.path.join(top_level, ".mobconf")):
    new_top_level = os.path.dirname(top_level)
    if new_top_level == top_level:
        print("fatal: could not find .mobconf")
        sys.exit(1)
    top_level = new_top_level

verbose_chdir(top_level)

mobconf = configparser.ConfigParser()
mobconf.read(".mobconf")
image_conf = mobconf['image'] or {}

image_url = image_conf['url'] if 'url' in image_conf else ''
if not image_url:
    raise Exception("Missing image.url in .mobconf")

tag = args.tag or image_conf['tag'] if 'tag' in image_conf else ''
if not tag:
    raise Exception("Pass a tag via --tag or image.tag in .mobconf")

image_and_tag = f'{image_url}:{tag}'

if not args.no_pull:
    maybe_run(["docker", "image", "pull", image_and_tag])

##
# docker run flags
##

# Compute mount_point, mount_from, and workdir
mount_point = "/tmp/mobilenode"
workdir = mount_point
mount_from = top_level

# docker-run parameters:
docker_run = ["docker",
              "run",
              # Remove container afterwards.
              "--rm",
              # Give docker its own subdirectory under target so it doesn't trash
              # the host target directory, but will still be cleaned up with it.
              "--env", "CARGO_TARGET_DIR=" + mount_point + "/target/docker",
              # This is required to run the fog unit tests
              "--env", "TEST_DATABASE_URL=postgres://localhost",
              # This is consumed by servers when running locally in the container
              "--env", "MC_CHAIN_ID=local",
              "--volume", mount_from + ":" + mount_point,
              "--workdir", workdir]

# the entrypoint script will pick up these vars and set up the user.
if not args.run_as_root:
    # figure out user and group
    uid = os.getuid()
    gid = os.getgid()
    username = getpass.getuser()
    groupname = grp.getgrgid(gid)[0]
    docker_run.extend([
            "--env", "EXTERNAL_UID=" + str(uid),
            "--env", "EXTERNAL_GID=" + str(gid),
            "--env", "EXTERNAL_USER=" + username,
            "--env", "EXTERNAL_GROUP=" + groupname
    ])

# Add /dev/isgx if HW mode, and it is available
# Mimicking bash checks `if [ -c /dev/isgx ]`
# The purpose of this check is that it is possible to build in HW mode, and even
# to run many tests, without installing sgx on the host and enabling this device in docker,
# but sometimes you really do need it, if you want to run consensus nodes etc.
# The friendliest thing seems to be, pass on the device if possible, maybe print a warning if not.
if args.hw:
    if pathlib.Path("/dev/isgx").is_char_device():
        docker_run.extend(["--device", "/dev/isgx"])
    else:
        eprint("Did not find /dev/isgx on the host, skipping")

# Add build environment
build_env = [
    "RUST_BACKTRACE=full",
    "SGX_MODE=HW" if args.hw else "SGX_MODE=SW",
    "IAS_MODE=PROD" if args.ias_prod else "IAS_MODE=DEV",
]

# Add GIT_COMMIT if present
git_commit = get_git_commit()
if git_commit:
    build_env.extend([f'GIT_COMMIT={git_commit}'])

for pair in build_env:
    docker_run.extend(["--env", pair])

# If running interactively (with a tty), get a tty in the container also
# This allows colored build logs when running locally without messing up logs in CI
if sys.stdout.isatty():
    docker_run.extend(["-t"])

ports = []
# in prompt, use -i to get user input
if args.action == "prompt":
    docker_run.extend(["-i"])

    # ports might be exposed to run clients or a local network
    ports = [
        "8080",
        "8081",
        "8443",
        "3223",
        "3225",
        "3226",
        "3228",
        "4444",
    ]

    # debug options allow us to attach gdb when debugging failing tests
    docker_run.extend([
        "--cap-add", "SYS_PTRACE",
    ])

ports.extend(args.expose or [])
for port in ports:
    docker_run.extend(["--expose", port])

for port in args.publish or []:
    docker_run.extend(["--publish", "{}:{}".format(port, port)])

# Map in the ssh directory, or the ssh-agent socket
if args.ssh_dir:
    eprint("SSH Agent authentication disabled. You will need to run the following commands to build: ")
    eprint('eval `ssh-agent`')
    eprint('ssh-add /root/.ssh/<your-ssh-private-key>')
    docker_run.extend(["--volume", os.path.expanduser("~/.ssh") + ":" + "/root/.ssh"])
elif args.ssh_agent:
    docker_run.extend(["--env", "SSH_AUTH_SOCK=/tmp/ssh_auth_sock"])
    docker_run.extend(["--volume", os.environ["SSH_AUTH_SOCK"] + ":" + "/tmp/ssh_auth_sock"])

# Name the container if a name was provided
if args.name:
    docker_run.extend(["--name", args.name])

# Add image name and command
docker_run.extend([image_and_tag])

docker_run.extend(["/bin/bash"])

if args.cmd:
    cmd = ' '.join(args.cmd)
    vprint('Running command in docker:', cmd)
    docker_run.extend(["-c", cmd])

try:
    maybe_run(docker_run)
except subprocess.CalledProcessError as exception:
    if args.cmd:
        # Make sure custom commands have their expected return codes
        sys.exit(exception.returncode)
    if args.action == 'prompt' and exception.returncode == 130:
        # This is a normal exit of prompt
        sys.exit(0)
    raise  # rethrow