cmd_capabilities.py: Port to new client channel helpers

#144

buildgrid/_app/commands/cmd_capabilities.py

@@ -14,12 +14,13 @@
 
 
 import sys
-from urllib.parse import urlparse
 
 import click
-import grpc
+from google.protobuf import json_format
 
+from buildgrid.client.authentication import setup_channel
 from buildgrid.client.capabilities import CapabilitiesInterface
+from buildgrid._exceptions import InvalidArgumentError
 
 from ..cli import pass_context
 
@@ -27,32 +28,30 @@ from ..cli import pass_context
 @click.command(name='capabilities', short_help="Capabilities service.")
 @click.option('--remote', type=click.STRING, default='http://localhost:50051', show_default=True,
               help="Remote execution server's URL (port defaults to 50051 if no specified).")
+@click.option('--auth-token', type=click.Path(exists=True, dir_okay=False), default=None,
+              help="Authorization token for the remote.")
 @click.option('--client-key', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Private client key for TLS (PEM-encoded)")
+              help="Private client key for TLS (PEM-encoded).")
 @click.option('--client-cert', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Public client certificate for TLS (PEM-encoded)")
+              help="Public client certificate for TLS (PEM-encoded).")
 @click.option('--server-cert', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Public server certificate for TLS (PEM-encoded)")
+              help="Public server certificate for TLS (PEM-encoded).")
 @click.option('--instance-name', type=click.STRING, default='main', show_default=True,
               help="Targeted farm instance name.")
 @pass_context
-def cli(context, remote, instance_name, client_key, client_cert, server_cert):
-    click.echo("Getting capabilities...")
-    url = urlparse(remote)
-
-    remote = '{}:{}'.format(url.hostname, url.port or 50051)
-    instance_name = instance_name
-
-    if url.scheme == 'http':
-        channel = grpc.insecure_channel(remote)
-    else:
-        credentials = context.load_client_credentials(client_key, client_cert, server_cert)
-        if not credentials:
-            click.echo("ERROR: no TLS keys were specified and no defaults could be found.", err=True)
-            sys.exit(-1)
-
-        channel = grpc.secure_channel(remote, credentials)
-
-    interface = CapabilitiesInterface(channel)
-    response = interface.get_capabilities(instance_name)
-    click.echo(response)
+def cli(context, remote, instance_name, auth_token, client_key, client_cert, server_cert):
+    """Entry point for the bgd-capabilities CLI command group."""
+    try:
+        context.channel, _ = setup_channel(remote, auth_token=auth_token,
+                                           client_key=client_key, client_cert=client_cert)
+
+    except InvalidArgumentError as e:
+        click.echo("Error: {}.".format(e), err=True)
+        sys.exit(-1)
+
+    context.instance_name = instance_name
+
+    interface = CapabilitiesInterface(context.channel)
+    response = interface.get_capabilities(context.instance_name)
+
+    click.echo(json_format.MessageToJson(response))