Merge pull request #714 from splitgraph/add-sgr-cloud-tunnel-CU-38f3r2e

Add sgr cloud tunnel command to start tunnel

Author: neumark

splitgraph/cloud/__init__.py

@@ -61,6 +61,8 @@
     INGESTION_JOB_STATUS,
     JOB_LOGS,
     PROFILE_UPSERT,
+    PROVISION_EPHEMERAL_TUNNEL,
+    PROVISION_REPOSITORY_TUNNEL,
     REPO_CONDITIONS,
     REPO_PARAMS,
     START_EXPORT,
@@ -1104,3 +1106,37 @@ def load_all_repositories(self, limit_to: List[str] = None) -> List[Repository]:
             parsed_external = ExternalResponse.from_response(external_r.json())
 
         return make_repositories(parsed_metadata, parsed_external)
+
+    def provision_repository_tunnel(self, namespace: str, repository: str) -> Tuple[str, str, int]:
+        response = self._gql(
+            {
+                "query": PROVISION_REPOSITORY_TUNNEL,
+                "operationName": "ProvisionRepositoryTunnel",
+                "variables": {"namespace": namespace, "repository": repository},
+            },
+            handle_errors=True,
+            anonymous_ok=False,
+        )
+        return (
+            response.json()["data"]["provisionRepositoryTunnel"]["secretToken"],
+            response.json()["data"]["provisionRepositoryTunnel"]["tunnelConnectHost"],
+            response.json()["data"]["provisionRepositoryTunnel"]["tunnelConnectPort"],
+        )
+
+    def provision_ephemeral_tunnel(self) -> Tuple[str, str, int, str, int]:
+        response = self._gql(
+            {
+                "query": PROVISION_EPHEMERAL_TUNNEL,
+                "operationName": "ProvisionEphemeralTunnel",
+                "variables": {},
+            },
+            handle_errors=True,
+            anonymous_ok=False,
+        )
+        return (
+            response.json()["data"]["provisionEphemeralTunnel"]["secretToken"],
+            response.json()["data"]["provisionEphemeralTunnel"]["tunnelConnectHost"],
+            response.json()["data"]["provisionEphemeralTunnel"]["tunnelConnectPort"],
+            response.json()["data"]["provisionEphemeralTunnel"]["privateAddressHost"],
+            response.json()["data"]["provisionEphemeralTunnel"]["privateAddressPort"],
+        )

splitgraph/cloud/queries.py

@@ -280,6 +280,26 @@
 }
 """
 
+PROVISION_REPOSITORY_TUNNEL = """mutation ProvisionRepositoryTunnel($namespace: String!, $repository: String!) {
+  provisionRepositoryTunnel(namespace:$namespace, repository:$repository) {
+    secretToken,
+    tunnelConnectHost,
+    tunnelConnectPort
+  }
+}
+"""
+
+PROVISION_EPHEMERAL_TUNNEL = """mutation ProvisionEphemeralTunnel {
+  provisionEphemeralTunnel {
+    secretToken,
+    tunnelConnectHost,
+    tunnelConnectPort,
+    privateAddressHost,
+    privateAddressPort
+  }
+}
+"""
+
 EXPORT_JOB_STATUS = """query ExportJobStatus($taskId: UUID!) {
   exportJobStatus(taskId: $taskId) {
     taskId

splitgraph/cloud/tunnel_client.py

@@ -0,0 +1,86 @@
+import os
+import subprocess
+from os import path
+from typing import Optional
+
+from click import echo
+
+RATHOLE_CLIENT_CONFIG_FILENAME = "rathole-client.toml"
+
+RATHOLE_CLIENT_CONFIG_TEMPLATE = """
+[client]
+remote_addr = "{tunnel_connect_address}"
+
+[client.transport]
+type = "tls"
+
+[client.transport.tls]
+{trusted_root_line}
+hostname = "{tls_hostname}"
+
+[client.services."{section_id}"]
+local_addr = "{local_address}"
+token = "{secret_token}"
+
+"""
+
+
+def get_rathole_client_config(
+    tunnel_connect_address: str,
+    tls_hostname: str,
+    local_address: str,
+    secret_token: str,
+    section_id: str,
+    trusted_root: Optional[str],
+) -> str:
+    trusted_root_line = f'trusted_root = "{trusted_root}"' if trusted_root else ""
+    return RATHOLE_CLIENT_CONFIG_TEMPLATE.format(
+        tunnel_connect_address=tunnel_connect_address,
+        tls_hostname=tls_hostname,
+        local_address=local_address,
+        secret_token=secret_token,
+        section_id=section_id,
+        trusted_root_line=trusted_root_line,
+    )
+
+
+def get_config_dir():
+    from splitgraph.config import CONFIG
+    from splitgraph.config.config import get_singleton
+
+    return os.path.dirname(get_singleton(CONFIG, "SG_CONFIG_FILE"))
+
+
+def get_rathole_client_binary_path():
+    return os.path.join(get_config_dir(), "rathole")
+
+
+def write_rathole_client_config(
+    section_id: str,
+    secret_token: str,
+    tunnel_connect_host: str,
+    tunnel_connect_port: int,
+    local_address: str,
+    tls_hostname: Optional[str],
+) -> str:
+    # If a specific root CA file is used (eg: for self-signed hosts), reference
+    # it in the rathole client config.
+    trusted_root = os.environ.get("REQUESTS_CA_BUNDLE") or os.environ.get("SSL_CERT_FILE")
+    rathole_client_config = get_rathole_client_config(
+        tunnel_connect_address=f"{tunnel_connect_host}:{tunnel_connect_port}",
+        tls_hostname=tls_hostname or tunnel_connect_host,
+        local_address=local_address,
+        secret_token=secret_token,
+        section_id=section_id,
+        trusted_root=trusted_root,
+    )
+    config_filename = path.join(get_config_dir(), RATHOLE_CLIENT_CONFIG_FILENAME)
+    with open(config_filename, "w") as f:
+        f.write(rathole_client_config)
+    return config_filename
+
+
+def launch_rathole_client(rathole_client_config_path: str) -> None:
+    echo("launching rathole client")
+    command = [get_rathole_client_binary_path(), "--client", rathole_client_config_path]
+    subprocess.check_call(command)

splitgraph/commandline/cloud.py

@@ -17,6 +17,10 @@
 
 from splitgraph.cloud.models import AddExternalRepositoryRequest, IntrospectionMode
 from splitgraph.cloud.project.models import Metadata, SplitgraphYAML
+from splitgraph.cloud.tunnel_client import (
+    launch_rathole_client,
+    write_rathole_client_config,
+)
 from splitgraph.commandline.common import (
     ImageType,
     RepositoryType,
@@ -1213,6 +1217,92 @@ def seed_c(remote, seed, github_repository, directory):
     click.echo(f"Splitgraph project generated in {os.path.abspath(directory)}.")
 
 
+def start_repository_tunnel(
+    remote: str, repository: "CoreRepository", external: "External"
+) -> None:
+    if not external.tunnel:
+        raise click.UsageError(
+            f"Repository {repository.namespace}/{repository.repository} is not tunneled"
+        )
+
+    # TODO: Get current version of rathole client for architecture.
+    # user must manually download rathole for now
+
+    from splitgraph.cloud import GQLAPIClient
+
+    client = GQLAPIClient(remote)
+    (secret_token, tunnel_connect_host, tunnel_connect_port) = client.provision_repository_tunnel(
+        repository.namespace, repository.repository
+    )
+
+    section_id = f"{repository.namespace}/{repository.repository}"
+    local_address = f"{external.params['host']}:{external.params['port']}"
+
+    rathole_client_config_path = write_rathole_client_config(
+        section_id,
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        local_address,
+        tunnel_connect_host,
+    )
+
+    launch_rathole_client(rathole_client_config_path)
+
+
+def start_ephemeral_tunnel(remote: str, local_address: str) -> None:
+    from splitgraph.cloud import GQLAPIClient
+
+    client = GQLAPIClient(remote)
+    (
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        private_address_host,
+        private_address_port,
+    ) = client.provision_ephemeral_tunnel()
+
+    rathole_client_config_path = write_rathole_client_config(
+        private_address_host,
+        secret_token,
+        tunnel_connect_host,
+        tunnel_connect_port,
+        local_address,
+        tunnel_connect_host,
+    )
+    click.echo(
+        f"To connect to {local_address} from Splitgraph, use the following connection parameters:\nHost: {private_address_host}\nPort: {private_address_port}"
+    )
+    launch_rathole_client(rathole_client_config_path)
+
+
+@click.command("tunnel")
+@click.option("--remote", default="data.splitgraph.com", help="Name of the remote registry to use.")
+@click.option(
+    "--repositories-file", "-f", default=["splitgraph.yml"], type=click.Path(), multiple=True
+)
+@click.argument("repository_or_local_address", type=str)
+def tunnel_c(remote: str, repositories_file: List[Path], repository_or_local_address: str):
+    """
+    Start the tunnel client to make tunneled external repo available.
+
+    This will load a splitgraph.yml file and tunnel the host:port address of the
+    external repository specified in the argument.
+    """
+
+    if "/" in repository_or_local_address:
+        repository: "CoreRepository" = RepositoryType(exists=False).convert(
+            repository_or_local_address, None, None
+        )
+        external = _get_external_from_yaml(repositories_file, repository)[0]
+        start_repository_tunnel(remote, repository, external)
+
+    elif ":" in repository_or_local_address:
+        start_ephemeral_tunnel(remote, repository_or_local_address)
+    else:
+        raise click.UsageError("Argument should be of the form namespace/repository or host:port")
+
+
 @click.group("cloud")
 def cloud_c():
     """Run actions on Splitgraph Cloud."""
@@ -1240,3 +1330,4 @@ def cloud_c():
 cloud_c.add_command(stub_c)
 cloud_c.add_command(validate_c)
 cloud_c.add_command(seed_c)
+cloud_c.add_command(tunnel_c)