Merge pull request #6575 from NHSDigital/add-mavis-server

Add mavis server

Author: misaka

.gitignore

@@ -90,3 +90,7 @@ dump.rdb
 # Data that gets downloaded by CLI tools
 db/data/dfe-schools.zip
 db/data/nhs-gp-practices.csv
+
+# Python
+__pycache__/
+*.pyc

.tool-versions

@@ -4,6 +4,7 @@ hk         1.37.0
 nodejs     22.15.0
 pkl        0.31.0
 postgres   17.2
+python     3.14.4
 redis      8.2.1
 ruby       4.0.2
 shellcheck 0.11.0

bin/mavis-server

@@ -0,0 +1,11 @@
+#/usr/bin/env sh
+
+set -eu
+
+bin_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
+root_dir=$(CDPATH= cd -- "$bin_dir/.." && pwd)
+
+export PYTHONPATH=$root_dir/python
+
+exec python -m mavis.server "$@"
+

docs/aws-setup.md

@@ -35,4 +35,4 @@ sso_registration_scopes = sso:account:access
 - Run `aws configure sso`. This will prompt you to log in to your AWS account and grant the necessary permissions for the CLI to access AWS services. When prompted for a region enter `eu-west-2` and for output format enter `json`.
 - Install the Session Manager plugin for the AWS CLI by following the instructions in the [AWS Systems Manager Session Manager documentation](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html).
 - Run `aws sso login` to log in to your AWS account and establish a session. This will allow you to access AWS resources using the CLI.
-- You should now be able to shell into a running service. The simplest way to do this is using the `script/shell.sh` script, e.g. `script/shell.sh qa`.
+- You should now be able to shell into a running service. The simplest way to do this is using the `bin/mavis-server shell` command, e.g. `bin/mavis-server shell qa`.

python/mavis/__init__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

python/mavis/server/__init__.py

@@ -0,0 +1,194 @@
+import json
+import subprocess
+
+from .helpers import run_command
+
+REGION = "eu-west-2"
+PRODUCTION_ENVS = {"production", "production-data-replication"}
+
+
+def cluster(env):
+    return f"mavis-{env}"
+
+
+def s3_bucket(env):
+    if env in PRODUCTION_ENVS:
+        return "mavis-filetransfer-production"
+    return "mavis-filetransfer-development"
+
+
+def ensure_authenticated(exit_without_login=False):
+    """Check AWS auth; attempt SSO login if needed."""
+    result = subprocess.run(
+        ["aws", "sts", "get-caller-identity"],
+        capture_output=True,
+    )
+    if result.returncode == 0:
+        return
+    if exit_without_login:
+        raise RuntimeError(
+            "Not authenticated with AWS. Run 'aws sso login' and try again."
+        )
+    print("Not authenticated with AWS. Attempting SSO login...")
+    login = subprocess.run(["aws", "sso", "login"])
+    if login.returncode != 0:
+        raise RuntimeError("AWS SSO login failed.")
+    recheck = subprocess.run(
+        ["aws", "sts", "get-caller-identity"],
+        capture_output=True,
+    )
+    if recheck.returncode != 0:
+        raise RuntimeError("Still not authenticated after SSO login.")
+
+
+def aws_json(*cmd):
+    """Run an AWS CLI command and return parsed JSON output."""
+    result = subprocess.run(["aws", *cmd], capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"aws {' '.join(cmd)}:\n{result.stderr.strip()}")
+    return json.loads(result.stdout)
+
+
+def resolve_task(env, task_id=None, task_ip=None, service=None):
+    """
+    Resolve to (short_task_id, container_name). Three mutually exclusive modes:
+
+    - task_id  — validate the specific task is running
+    - task_ip  — search all running tasks in the cluster for a matching IP
+    - service  — return the first running task in the service; defaults to
+                 mavis-{env}-ops, or mavis-{env}-web for data-replication envs
+    """
+    cl = cluster(env)
+
+    if task_id:
+        tasks = aws_json(
+            "ecs",
+            "describe-tasks",
+            "--region",
+            REGION,
+            "--cluster",
+            cl,
+            "--tasks",
+            task_id,
+        ).get("tasks", [])
+        if not tasks or tasks[0]["lastStatus"] != "RUNNING":
+            raise RuntimeError(f"Task {task_id} is not running in cluster {cl}")
+        return task_id, _application_container(tasks[0])
+
+    if task_ip:
+        task_arns = aws_json(
+            "ecs",
+            "list-tasks",
+            "--region",
+            REGION,
+            "--cluster",
+            cl,
+            "--desired-status",
+            "RUNNING",
+        ).get("taskArns", [])
+        if not task_arns:
+            raise RuntimeError(f"No running tasks found in cluster {cl}")
+        tasks = aws_json(
+            "ecs",
+            "describe-tasks",
+            "--region",
+            REGION,
+            "--cluster",
+            cl,
+            "--tasks",
+            *task_arns,
+        ).get("tasks", [])
+        for task in tasks:
+            if _task_private_ip(task) == task_ip:
+                return _short_id(task), _application_container(task)
+        raise RuntimeError(f"No running task with IP {task_ip} found in cluster {cl}")
+
+    if not service:
+        service = _default_service(env)
+
+    task_arns = aws_json(
+        "ecs",
+        "list-tasks",
+        "--region",
+        REGION,
+        "--cluster",
+        cl,
+        "--service-name",
+        service,
+        "--desired-status",
+        "RUNNING",
+    ).get("taskArns", [])
+    if not task_arns:
+        raise RuntimeError(f"No running tasks found in service {service}")
+    tasks = aws_json(
+        "ecs",
+        "describe-tasks",
+        "--region",
+        REGION,
+        "--cluster",
+        cl,
+        "--tasks",
+        *task_arns,
+    ).get("tasks", [])
+    for task in tasks:
+        container = _application_container(task)
+        if container:
+            return _short_id(task), container
+    raise RuntimeError(
+        f"No running tasks with an application container found in service {service}"
+    )
+
+
+def run_remote_command(
+    env, task_id, remote_command, container=None, replace_process=False
+):
+    """Execute a command in an ECS task, returning the exit code."""
+    command = [
+        "aws",
+        "ecs",
+        "execute-command",
+        "--region",
+        REGION,
+        "--cluster",
+        cluster(env),
+        "--task",
+        task_id,
+        "--command",
+        remote_command,
+        "--interactive",
+    ]
+    if container:
+        command += ["--container", container]
+    return run_command(command, replace_process=replace_process)
+
+
+# --- private helpers ---
+
+
+def _default_service(env):
+    if env.endswith("data-replication"):
+        return f"mavis-{env}"
+    return f"mavis-{env}-ops"
+
+
+def _short_id(task):
+    return task["taskArn"].split("/")[-1]
+
+
+def _application_container(task):
+    for c in task.get("containers", []):
+        if (
+            c.get("name") == "application"
+            and c.get("lastStatus") == "RUNNING"
+            and c.get("runtimeId")
+        ):
+            return c["name"]
+    return None
+
+
+def _task_private_ip(task):
+    for attachment in task.get("attachments", []):
+        for detail in attachment.get("details", []):
+            if detail.get("name") == "privateIPv4Address":
+                return detail.get("value")
+    return None

python/mavis/server/__main__.py

@@ -0,0 +1,19 @@
+import argparse
+
+from . import get_file, put_file, shell
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="mavis-server",
+        description="MAVIS server management CLI",
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    get_file.register(subparsers)
+    put_file.register(subparsers)
+    shell.register(subparsers)
+
+    args = parser.parse_args()
+    # TODO: Clean this error reporting up
+    args.func(args)

python/mavis/server/aws.py

@@ -0,0 +1,96 @@
+import os
+import secrets
+import sys
+
+from . import aws
+from .helpers import confirm_production, run_command
+
+
+def register(subparsers):
+    parser = subparsers.add_parser(
+        "get-file",
+        help="Download a file from an ECS container to local",
+        description=(
+            "Download a file from inside an ECS container to a local path, "
+            "using S3 as an intermediary. The S3 object is always cleaned up."
+        ),
+    )
+    parser.add_argument("env", help="Environment name (cluster will be mavis-ENV)")
+    parser.add_argument("remote_path", help="Path of the file inside the container")
+    parser.add_argument(
+        "local_path",
+        nargs="?",
+        default=None,
+        help="Local destination (file or directory). Defaults to tmp in the project root.",
+    )
+    parser.add_argument("--service", help="Override the ECS service name")
+    parser.add_argument(
+        "--task-id", dest="task_id", help="Connect to a specific task by ID"
+    )
+    parser.add_argument(
+        "--task-ip",
+        dest="task_ip",
+        help="Connect to a task by its private IPv4 address",
+    )
+    parser.add_argument(
+        "-x",
+        "--exit-without-login",
+        dest="exit_without_login",
+        action="store_true",
+        help="Exit instead of prompting for AWS SSO login",
+    )
+    parser.set_defaults(func=run)
+
+
+def run(args):
+    env = args.env
+
+    confirm_production(env)
+    aws.ensure_authenticated(exit_without_login=args.exit_without_login)
+
+    task_id, container = aws.resolve_task(
+        env, task_id=args.task_id, task_ip=args.task_ip, service=args.service
+    )
+    bucket = aws.s3_bucket(env)
+    key = f"temp-{secrets.token_hex(8)}"
+    s3_uri = f"s3://{bucket}/{key}"
+
+    local_dest = _local_destination(args.remote_path, args.local_path)
+
+    try:
+        upload_result = aws.run_remote_command(
+            env,
+            task_id,
+            f"aws s3 cp {args.remote_path} {s3_uri} --region {aws.REGION}",
+            container=container,
+        )
+        if not upload_result:
+            sys.exit("Error: Failed to copy file from container to S3")
+
+        download_result = run_command(
+            ["aws", "s3", "cp", s3_uri, local_dest, "--region", aws.REGION]
+        )
+        if not download_result:
+            sys.exit("Error: Download from S3 failed with code")
+    finally:
+        run_command(
+            ["aws", "s3", "rm", s3_uri, "--region", aws.REGION],
+        )
+
+    print(f"File successfully downloaded to {local_dest}")
+
+
+def _local_destination(remote_path, local_path):
+    """
+    Resolve the local download destination.
+
+    If local_path is given and is an existing directory, save as
+    <local_path>/<basename of remote_path>. If local_path is a file path
+    (or doesn't exist yet), use it as-is. Defaults to ./<basename>.
+    """
+    filename = os.path.basename(remote_path.rstrip("/"))
+    if local_path is None:
+        return os.path.join(".", filename)
+    if os.path.isdir(local_path):
+        return os.path.join(local_path, filename)
+    return local_path