diff --git a/morphcloud/cli_plugins/__init__.py b/morphcloud/cli_plugins/__init__.py new file mode 100644 index 0000000..27d8979 --- /dev/null +++ b/morphcloud/cli_plugins/__init__.py @@ -0,0 +1,5 @@ +"""morphcloud CLI plugin modules. + +Plugins are loaded via the `morphcloud.cli_plugins` entry point group. +""" + diff --git a/morphcloud/cli_plugins/aws_sim.py b/morphcloud/cli_plugins/aws_sim.py new file mode 100644 index 0000000..d361e06 --- /dev/null +++ b/morphcloud/cli_plugins/aws_sim.py @@ -0,0 +1,438 @@ +import json +import os +import pathlib +import subprocess +from typing import Any + +import click +import httpx + + +DEFAULT_SIM_AWS_BASE_URL = "https://sim-aws.svc.cloud.morph.so" +DEFAULT_CONNECTOR_IMAGE = "ghcr.io/morph-labs/sim-aws-connector:latest" +DEFAULT_CONNECT_BUNDLE_FILENAME = "aws-sim-connect-bundle.json" +SRC_VALID_MARK_SYSCTL = "net.ipv4.conf.all.src_valid_mark=1" +DEFAULT_CONNECTOR_CONTAINER_NAME = "sim-aws-connector" +DEFAULT_CONNECT_HELPERS_ENV = "aws-env.sh" +DEFAULT_CONNECT_HELPERS_WRAPPER = "aws" + + +def _get_sim_aws_base_url() -> str: + return os.environ.get("SIM_AWS_BASE_URL", DEFAULT_SIM_AWS_BASE_URL).rstrip("/") + + +def _get_morph_api_key() -> str: + key = os.environ.get("MORPH_API_KEY") + if not key: + raise click.ClickException( + "Error: MORPH_API_KEY environment variable is not set. " + "Set it (export MORPH_API_KEY='...') and retry." + ) + return key + + +def _http_client() -> httpx.Client: + return httpx.Client( + base_url=_get_sim_aws_base_url(), + headers={"Authorization": f"Bearer {_get_morph_api_key()}"}, + timeout=httpx.Timeout(30.0), + ) + + +def _raise_for_status(resp: httpx.Response) -> None: + try: + resp.raise_for_status() + except httpx.HTTPStatusError as e: + body = "" + try: + body = resp.text + except Exception: + body = "" + msg = f"Sim-AWS request failed: {resp.status_code} {resp.reason_phrase}" + if body: + msg = f"{msg}\n{body}" + raise click.ClickException(msg) from e + + +def _request_json(method: str, path: str, *, json_body: Any | None = None) -> Any: + with _http_client() as client: + resp = client.request(method, path, json=json_body) + _raise_for_status(resp) + if resp.status_code == 204: + return None + return resp.json() + + +def _ensure_group(cli_group: click.Group, name: str, help_text: str) -> click.Group: + existing = cli_group.commands.get(name) + if isinstance(existing, click.Group): + return existing + if existing is not None: + raise click.ClickException( + f"Cannot install aws-sim plugin: command '{name}' already exists and is not a group." + ) + new_group = click.Group(name=name, help=help_text) + cli_group.add_command(new_group) + return new_group + + +def _print_json(data: Any) -> None: + click.echo(json.dumps(data, indent=2, sort_keys=True)) + + +def _docker_run_template(bundle_path: str) -> str: + bundle_path_abs = str(pathlib.Path(bundle_path).expanduser().resolve()) + image = os.environ.get("AWS_SIM_CONNECTOR_IMAGE", DEFAULT_CONNECTOR_IMAGE) + bundle_container_path = "/run/connect-bundle.json" + return ( + "docker run --rm -it " + "--cap-add=NET_ADMIN " + "--device /dev/net/tun " + f"--sysctl {SRC_VALID_MARK_SYSCTL} " + "-e MORPH_API_KEY " + f"-v {bundle_path_abs}:{bundle_container_path}:ro " + f"{image} " + "bash" + ) + + +def _docker_run_detached_template(bundle_path: str, *, container_name: str = "sim-aws-connector") -> str: + bundle_path_abs = str(pathlib.Path(bundle_path).expanduser().resolve()) + image = os.environ.get("AWS_SIM_CONNECTOR_IMAGE", DEFAULT_CONNECTOR_IMAGE) + bundle_container_path = "/run/connect-bundle.json" + return ( + f"docker run -d --name {container_name} " + "--cap-add=NET_ADMIN " + "--device /dev/net/tun " + f"--sysctl {SRC_VALID_MARK_SYSCTL} " + "-e MORPH_API_KEY " + "-e AWS_PAGER= " + "-e AWS_EC2_METADATA_DISABLED=true " + "-e AWS_ACCESS_KEY_ID=test " + "-e AWS_SECRET_ACCESS_KEY=test " + f"-v {bundle_path_abs}:{bundle_container_path}:ro " + f"{image} " + "sleep infinity" + ) + +def _default_region_from_bundle(bundle: Any) -> str: + aws = bundle.get("aws") if isinstance(bundle, dict) else None + if not isinstance(aws, dict): + return "" + regions = aws.get("regions") + if not isinstance(regions, list) or not regions: + return "" + return str(regions[0] or "").strip() + + +def _emit_connect_helpers( + *, + dir_path: pathlib.Path, + container_name: str, +) -> tuple[pathlib.Path, pathlib.Path]: + dir_path.mkdir(parents=True, exist_ok=True) + + env_path = dir_path / DEFAULT_CONNECT_HELPERS_ENV + wrapper_path = dir_path / DEFAULT_CONNECT_HELPERS_WRAPPER + + env_path.write_text( + "\n".join( + [ + "# Source this file to make `aws` call the local ./aws wrapper (relative to this file).", + '_simaws_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"', + f'export AWS_SIM_CONNECTOR_CONTAINER="${{AWS_SIM_CONNECTOR_CONTAINER:-{container_name}}}"', + 'aws() { "$_simaws_dir/aws" "$@"; }', + ] + ), + encoding="utf-8", + ) + + wrapper_path.write_text( + "\n".join( + [ + "#!/usr/bin/env bash", + "set -euo pipefail", + "", + 'CONTAINER="${AWS_SIM_CONNECTOR_CONTAINER:-sim-aws-connector}"', + "", + 'if ! docker ps --format "{{.Names}}" | grep -qx "$CONTAINER"; then', + ' echo "ERROR: connector container \'$CONTAINER\' is not running." >&2', + " exit 1", + "fi", + "", + "tty=()", + '[[ -t 0 && -t 1 ]] && tty=(-t)', + "", + 'exec docker exec -i "${tty[@]}" "$CONTAINER" bash -lc \'', + " set -euo pipefail", + " # Import AWS_* (and CA bundle vars) from PID1 env (set by connector entrypoint).", + " while IFS= read -r kv; do", + " case \"$kv\" in", + " AWS_*=*|SSL_CERT_FILE=*|REQUESTS_CA_BUNDLE=*|CURL_CA_BUNDLE=*)", + " export \"$kv\"", + " ;;", + " esac", + " done < <(tr \"\\0\" \"\\n\" list[str]: + items: list[str] = [] + for raw in values: + for part in raw.split(","): + part = part.strip() + if part: + items.append(part) + return items + + +def _default_regions() -> list[str]: + raw = os.environ.get("SIM_AWS_REGIONS", "").strip() + if raw: + return _split_csv_args((raw,)) + return ["us-east-1"] + + +def _default_services() -> list[str]: + raw = os.environ.get("SIM_AWS_SERVICES", "").strip() + if raw: + return _split_csv_args((raw,)) + return ["s3", "ec2"] + + +def load(cli_group: click.Group) -> None: + """ + Entry point for the `morphcloud.cli_plugins` group. + + Registers: `morphcloud env aws-sim ...` + """ + + env_group = _ensure_group(cli_group, "env", "Manage external environments.") + aws_sim_group = _ensure_group(env_group, "aws-sim", "Manage Sim-AWS environments.") + + @aws_sim_group.command(name="create") + @click.option( + "--region", + "regions", + multiple=True, + help="Region(s); repeatable and/or comma-separated.", + ) + @click.option( + "--service", + "services", + multiple=True, + help="Service(s); repeatable and/or comma-separated.", + ) + @click.option( + "--ttl-seconds", + "ttl_seconds", + type=int, + default=None, + help="Environment TTL in seconds.", + ) + @click.option( + "--name", + "name", + type=str, + default=None, + help="Optional environment name.", + ) + def create(regions: tuple[str, ...], services: tuple[str, ...], ttl_seconds: int | None, name: str | None) -> None: + """Create a Sim-AWS environment.""" + create_body: dict[str, Any] = { + "regions": _split_csv_args(regions) or _default_regions(), + "services": _split_csv_args(services) or _default_services(), + } + if ttl_seconds is not None: + create_body["ttl_seconds"] = ttl_seconds + if name: + create_body["name"] = name + _print_json(_request_json("POST", "/v1/envs", json_body=create_body)) + + @aws_sim_group.command(name="list") + def list_() -> None: + """List Sim-AWS environments.""" + _print_json(_request_json("GET", "/v1/envs")) + + @aws_sim_group.command(name="get") + @click.argument("env_id") + def get(env_id: str) -> None: + """Get a Sim-AWS environment.""" + _print_json(_request_json("GET", f"/v1/envs/{env_id}")) + + @aws_sim_group.command(name="start") + @click.argument("env_id") + def start(env_id: str) -> None: + """Start a Sim-AWS environment.""" + _print_json(_request_json("POST", f"/v1/envs/{env_id}/start")) + + @aws_sim_group.command(name="pause") + @click.argument("env_id") + def pause(env_id: str) -> None: + """Pause a Sim-AWS environment.""" + _print_json(_request_json("POST", f"/v1/envs/{env_id}/pause")) + + @aws_sim_group.command(name="snapshot") + @click.argument("env_id") + def snapshot(env_id: str) -> None: + """Snapshot a Sim-AWS environment.""" + _print_json(_request_json("POST", f"/v1/envs/{env_id}/snapshot")) + + @aws_sim_group.command(name="restore") + @click.argument("env_id") + @click.argument("snapshot_id") + def restore(env_id: str, snapshot_id: str) -> None: + """Restore a Sim-AWS environment from a snapshot_id.""" + _print_json(_request_json("POST", f"/v1/envs/{env_id}/restore", json_body={"snapshot_id": snapshot_id})) + + @aws_sim_group.command(name="delete") + @click.argument("env_id") + def delete(env_id: str) -> None: + """Delete a Sim-AWS environment.""" + _print_json(_request_json("DELETE", f"/v1/envs/{env_id}")) + + @aws_sim_group.command(name="connect") + @click.argument("env_id") + @click.option( + "--output", + "output_path", + default=None, + show_default=False, + help="Optional: write the connect bundle to this path (sensitive).", + ) + @click.option( + "--container-name", + default=DEFAULT_CONNECTOR_CONTAINER_NAME, + show_default=True, + help="Name for the detached connector container.", + ) + @click.option( + "--replace", + is_flag=True, + help="If set, delete any existing container with --container-name before starting.", + ) + @click.option( + "--no-run", + is_flag=True, + help="If set, do not start the detached connector container; only print commands and emit helper scripts.", + ) + @click.option( + "--emit-dir", + default=".", + show_default=True, + help="Directory to write helper scripts (aws-env.sh + aws wrapper).", + ) + def connect( + env_id: str, + output_path: str | None, + container_name: str, + replace: bool, + no_run: bool, + emit_dir: str, + ) -> None: + """ + Fetch a connect bundle and start a detached connector container. + + By default this does NOT write the bundle to disk. Instead, it passes the bundle to the + connector via the CONNECT_BUNDLE_JSON environment variable. + + If you want a bundle file for debugging, pass --output. + """ + + bundle = _request_json("POST", f"/v1/envs/{env_id}/connect") + + output = None + if output_path: + output = pathlib.Path(output_path).expanduser() + output.write_text(json.dumps(bundle, indent=2) + "\n", encoding="utf-8") + try: + os.chmod(output, 0o600) + except Exception: + pass + click.echo(f"Wrote connect bundle to: {output}") + + default_region = _default_region_from_bundle(bundle) + + image = os.environ.get("AWS_SIM_CONNECTOR_IMAGE", DEFAULT_CONNECTOR_IMAGE) + requires_morph_api_key = bool(bundle.get("auth")) if isinstance(bundle, dict) else True + + # Always include dummy creds + pager off for smoother UX. + run_args: list[str] = [ + "docker", + "run", + "-d", + "--name", + container_name, + "--cap-add=NET_ADMIN", + "--device", + "/dev/net/tun", + "--sysctl", + SRC_VALID_MARK_SYSCTL, + "-e", + "CONNECT_BUNDLE_JSON", + "-e", + "AWS_PAGER=", + "-e", + "AWS_EC2_METADATA_DISABLED=true", + "-e", + "AWS_ACCESS_KEY_ID=test", + "-e", + "AWS_SECRET_ACCESS_KEY=test", + ] + if requires_morph_api_key: + run_args += ["-e", "MORPH_API_KEY"] + if default_region: + run_args += [ + "-e", + f"AWS_REGION={default_region}", + "-e", + f"AWS_DEFAULT_REGION={default_region}", + ] + run_args += [ + image, + "sleep", + "infinity", + ] + + helpers_dir = pathlib.Path(emit_dir).expanduser().resolve() + env_path, wrapper_path = _emit_connect_helpers( + dir_path=helpers_dir, + container_name=container_name, + ) + click.echo(f"Wrote helper scripts: {env_path} and {wrapper_path}") + click.echo(f"To enable `aws ...` in your current shell: source {env_path}") + + if no_run: + return + + if replace: + subprocess.run(["docker", "rm", "-f", container_name], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + try: + env = dict(os.environ) + # Compact JSON to avoid newlines; passed to docker via `-e CONNECT_BUNDLE_JSON`. + env["CONNECT_BUNDLE_JSON"] = json.dumps(bundle, separators=(",", ":"), sort_keys=True) + p = subprocess.run(run_args, text=True, capture_output=True, env=env) + except FileNotFoundError as e: + raise click.ClickException("docker is required but was not found on PATH") from e + + if p.returncode != 0: + msg = (p.stderr or p.stdout or "").strip() + raise click.ClickException(f"Failed to start connector container (docker run exit {p.returncode}):\n{msg}") + + container_id = (p.stdout or "").strip() + if container_id: + click.echo(f"Started connector container: {container_name} ({container_id[:12]})") diff --git a/pyproject.toml b/pyproject.toml index ee8be6c..644c14a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,9 @@ dev = [ [project.scripts] morphcloud = "morphcloud.cli:cli" +[project.entry-points."morphcloud.cli_plugins"] +aws-sim = "morphcloud.cli_plugins.aws_sim:load" + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/tests/unit/test_cli_env_aws_sim.py b/tests/unit/test_cli_env_aws_sim.py new file mode 100644 index 0000000..c43495e --- /dev/null +++ b/tests/unit/test_cli_env_aws_sim.py @@ -0,0 +1,300 @@ +import importlib +import sys + +from click.testing import CliRunner + + +def _import_cli_with_aws_sim_plugin(monkeypatch): + import importlib.metadata + + entry_point = importlib.metadata.EntryPoint( + name="aws-sim", + value="morphcloud.cli_plugins.aws_sim:load", + group="morphcloud.cli_plugins", + ) + + def fake_entry_points(*, group=None, **kwargs): + if group == "morphcloud.cli_plugins": + return [entry_point] + return [] + + monkeypatch.setattr(importlib.metadata, "entry_points", fake_entry_points) + sys.modules.pop("morphcloud.cli", None) + return importlib.import_module("morphcloud.cli") + + +def test_env_aws_sim_plugin_registers_commands(monkeypatch): + cli_mod = _import_cli_with_aws_sim_plugin(monkeypatch) + + runner = CliRunner() + result = runner.invoke(cli_mod.cli, ["env", "aws-sim", "--help"]) + assert result.exit_code == 0, result.output + + +def test_env_aws_sim_list_wires_httpx(monkeypatch): + cli_mod = _import_cli_with_aws_sim_plugin(monkeypatch) + + import morphcloud.cli_plugins.aws_sim as aws_sim_mod + + calls = [] + + class StubResponse: + status_code = 200 + reason_phrase = "OK" + text = "" + + def raise_for_status(self): + return None + + def json(self): + return {"ok": True} + + class StubClient: + def __init__(self, *, base_url, headers, timeout): + calls.append({"base_url": str(base_url), "headers": dict(headers), "timeout": timeout}) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def request(self, method, url, json=None): + calls[-1].update({"method": method, "url": url, "json": json}) + return StubResponse() + + monkeypatch.setattr(aws_sim_mod.httpx, "Client", StubClient) + monkeypatch.setenv("SIM_AWS_BASE_URL", "https://example.test") + monkeypatch.setenv("MORPH_API_KEY", "k_test_123") + + runner = CliRunner() + result = runner.invoke(cli_mod.cli, ["env", "aws-sim", "list"]) + assert result.exit_code == 0, result.output + assert calls == [ + { + "base_url": "https://example.test", + "headers": {"Authorization": "Bearer k_test_123"}, + "timeout": calls[0]["timeout"], + "method": "GET", + "url": "/v1/envs", + "json": None, + } + ] + assert "k_test_123" not in result.output + + +def test_env_aws_sim_create_sends_body(monkeypatch): + cli_mod = _import_cli_with_aws_sim_plugin(monkeypatch) + + import morphcloud.cli_plugins.aws_sim as aws_sim_mod + + calls = [] + + class StubResponse: + status_code = 200 + reason_phrase = "OK" + text = "" + + def raise_for_status(self): + return None + + def json(self): + return {"env_id": "awsenv_test"} + + class StubClient: + def __init__(self, *, base_url, headers, timeout): + calls.append({"base_url": str(base_url), "headers": dict(headers), "timeout": timeout}) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def request(self, method, url, json=None): + calls[-1].update({"method": method, "url": url, "json": json}) + return StubResponse() + + monkeypatch.setattr(aws_sim_mod.httpx, "Client", StubClient) + monkeypatch.setenv("SIM_AWS_BASE_URL", "https://example.test") + monkeypatch.setenv("MORPH_API_KEY", "k_test_123") + + runner = CliRunner() + result = runner.invoke( + cli_mod.cli, + [ + "env", + "aws-sim", + "create", + "--region", + "us-east-1", + "--service", + "s3", + "--ttl-seconds", + "3600", + "--name", + "test-env", + ], + ) + assert result.exit_code == 0, result.output + assert calls == [ + { + "base_url": "https://example.test", + "headers": {"Authorization": "Bearer k_test_123"}, + "timeout": calls[0]["timeout"], + "method": "POST", + "url": "/v1/envs", + "json": { + "name": "test-env", + "regions": ["us-east-1"], + "services": ["s3"], + "ttl_seconds": 3600, + }, + } + ] + assert "k_test_123" not in result.output + + +def test_env_aws_sim_restore_sends_snapshot_id(monkeypatch): + cli_mod = _import_cli_with_aws_sim_plugin(monkeypatch) + + import morphcloud.cli_plugins.aws_sim as aws_sim_mod + + calls = [] + + class StubResponse: + status_code = 200 + reason_phrase = "OK" + text = "" + + def raise_for_status(self): + return None + + def json(self): + return {"ok": True} + + class StubClient: + def __init__(self, *, base_url, headers, timeout): + calls.append({"base_url": str(base_url), "headers": dict(headers), "timeout": timeout}) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def request(self, method, url, json=None): + calls[-1].update({"method": method, "url": url, "json": json}) + return StubResponse() + + monkeypatch.setattr(aws_sim_mod.httpx, "Client", StubClient) + monkeypatch.setenv("SIM_AWS_BASE_URL", "https://example.test") + monkeypatch.setenv("MORPH_API_KEY", "k_test_123") + + runner = CliRunner() + result = runner.invoke( + cli_mod.cli, + ["env", "aws-sim", "restore", "awsenv_test", "awssnap_test"], + ) + assert result.exit_code == 0, result.output + assert calls == [ + { + "base_url": "https://example.test", + "headers": {"Authorization": "Bearer k_test_123"}, + "timeout": calls[0]["timeout"], + "method": "POST", + "url": "/v1/envs/awsenv_test/restore", + "json": {"snapshot_id": "awssnap_test"}, + } + ] + assert "k_test_123" not in result.output + + +def test_env_aws_sim_connect_writes_bundle_and_prints_docker_cmd(monkeypatch): + cli_mod = _import_cli_with_aws_sim_plugin(monkeypatch) + + import morphcloud.cli_plugins.aws_sim as aws_sim_mod + + connect_calls = [] + + class StubResponse: + status_code = 200 + reason_phrase = "OK" + text = "" + + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self._payload + + class StubClient: + def __init__(self, *, base_url, headers, timeout): + self._headers = dict(headers) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def request(self, method, url, json=None): + connect_calls.append({"method": method, "url": url, "headers": self._headers, "json": json}) + return StubResponse( + { + "version": "v1", + "env_id": "awsenv_test", + "instance_id": "morphvm_test", + "tunnel_ws_url": "wss://example.test/tunnel", + "wg": { + "client_address": "10.0.0.2/32", + "client_private_key": "base64-private", + "server_public_key": "base64-pub", + "allowed_ips": ["10.0.0.0/24"], + "endpoint_host": "127.0.0.1", + "endpoint_port": 51820, + "mtu": 1280, + "persistent_keepalive": 25, + }, + "dns": {"nameserver": "10.0.0.1"}, + "tls": {"ca_cert_pem": "pem", "ca_fingerprint_sha256": "fp"}, + "aws": {"gateway_ip": "10.0.0.3", "regions": ["us-east-1"]}, + "auth": { + "mode": "morph_api_key_bearer", + "header_name": "Authorization", + "header_value_template": "Bearer ${MORPH_API_KEY}", + }, + "notes": [], + } + ) + + monkeypatch.setattr(aws_sim_mod.httpx, "Client", StubClient) + monkeypatch.setenv("SIM_AWS_BASE_URL", "https://example.test") + monkeypatch.setenv("MORPH_API_KEY", "k_test_123") + + runner = CliRunner() + with runner.isolated_filesystem(): + result = runner.invoke( + cli_mod.cli, + ["env", "aws-sim", "connect", "awsenv_test", "--output", "bundle.json"], + ) + assert result.exit_code == 0, result.output + assert connect_calls == [ + { + "method": "POST", + "url": "/v1/envs/awsenv_test/connect", + "headers": {"Authorization": "Bearer k_test_123"}, + "json": None, + } + ] + assert "bundle.json" in result.output + assert "docker run" in result.output + assert "-e MORPH_API_KEY" in result.output + assert "--cap-add=NET_ADMIN" in result.output + assert "--device /dev/net/tun" in result.output + assert aws_sim_mod.SRC_VALID_MARK_SYSCTL in result.output + assert "k_test_123" not in result.output +