cmd_bot.py: Port to new client channel helpers

#144

buildgrid/_app/commands/cmd_bot.py

@@ -22,16 +22,15 @@ Create a bot interface and request work
 
 from pathlib import Path, PurePath
 import sys
-from urllib.parse import urlparse
 
 import click
-import grpc
 
 from buildgrid.bot import bot, interface, session
 from buildgrid.bot.hardware.interface import HardwareInterface
 from buildgrid.bot.hardware.device import Device
 from buildgrid.bot.hardware.worker import Worker
-
+from buildgrid.client.authentication import setup_channel
+from buildgrid._exceptions import InvalidArgumentError
 
 from ..bots import buildbox, dummy, host
 from ..cli import pass_context, setup_logging
@@ -40,20 +39,22 @@ from ..cli import pass_context, setup_logging
 @click.group(name='bot', short_help="Create and register bot clients.")
 @click.option('--remote', type=click.STRING, default='http://localhost:50051', show_default=True,
               help="Remote execution server's URL (port defaults to 50051 if not 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('--remote-cas', type=click.STRING, default=None, show_default=True,
               help="Remote CAS server's URL (port defaults to 11001 if not specified).")
 @click.option('--cas-client-key', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Private CAS client key for TLS (PEM-encoded)")
+              help="Private CAS client key for TLS (PEM-encoded).")
 @click.option('--cas-client-cert', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Public CAS client certificate for TLS (PEM-encoded)")
+              help="Public CAS client certificate for TLS (PEM-encoded).")
 @click.option('--cas-server-cert', type=click.Path(exists=True, dir_okay=False), default=None,
-              help="Public CAS server certificate for TLS (PEM-encoded)")
+              help="Public CAS server certificate for TLS (PEM-encoded).")
 @click.option('--update-period', type=click.FLOAT, default=0.5, show_default=True,
               help="Time period for bot updates to the server in seconds.")
 @click.option('--parent', type=click.STRING, default='main', show_default=True,
@@ -61,69 +62,31 @@ from ..cli import pass_context, setup_logging
 @click.option('-v', '--verbose', count=True,
               help='Increase log verbosity level.')
 @pass_context
-def cli(context, parent, update_period, remote, client_key, client_cert, server_cert,
+def cli(context, parent, update_period, remote, auth_token, client_key, client_cert, server_cert,
         remote_cas, cas_client_key, cas_client_cert, cas_server_cert, verbose):
     setup_logging(verbosity=verbose)
     # Setup the remote execution server channel:
-    url = urlparse(remote)
-
-    context.remote = '{}:{}'.format(url.hostname, url.port or 50051)
-    context.remote_url = remote
-    context.update_period = update_period
-    context.parent = parent
-
-    if url.scheme == 'http':
-        context.channel = grpc.insecure_channel(context.remote)
-
-        context.client_key = None
-        context.client_cert = None
-        context.server_cert = None
-    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)
-
-        context.channel = grpc.secure_channel(context.remote, credentials)
-
-        context.client_key = credentials.client_key
-        context.client_cert = credentials.client_cert
-        context.server_cert = credentials.server_cert
-
-    # Setup the remote CAS server channel, if separated:
-    if remote_cas is not None and remote_cas != remote:
-        cas_url = urlparse(remote_cas)
-
-        context.remote_cas = '{}:{}'.format(cas_url.hostname, cas_url.port or 11001)
-        context.remote_cas_url = remote_cas
+    try:
+        context.channel, details = setup_channel(remote, auth_token=auth_token, server_cert=server_cert,
+                                                 client_key=client_key, client_cert=client_cert)
 
-        if cas_url.scheme == 'http':
-            context.cas_channel = grpc.insecure_channel(context.remote_cas)
+        if remote_cas is not None and remote_cas != remote:
+            context.cas_channel, details = setup_channel(remote_cas, server_cert=cas_server_cert,
+                                                         client_key=cas_client_key, client_cert=cas_client_cert)
+            context.remote_cas_url = remote_cas
 
-            context.cas_client_key = None
-            context.cas_client_cert = None
-            context.cas_server_cert = None
         else:
-            cas_credentials = context.load_client_credentials(cas_client_key, cas_client_cert, cas_server_cert)
-            if not cas_credentials:
-                click.echo("ERROR: no TLS keys were specified and no defaults could be found.", err=True)
-                sys.exit(-1)
-
-            context.cas_channel = grpc.secure_channel(context.remote_cas, cas_credentials)
+            context.cas_channel = context.channel
+            context.remote_cas_url = remote
 
-            context.cas_client_key = cas_credentials.client_key
-            context.cas_client_cert = cas_credentials.client_cert
-            context.cas_server_cert = cas_credentials.server_cert
+        context.cas_client_key, context.cas_client_cert, context.cas_server_cert = details
 
-    else:
-        context.remote_cas = context.remote
-        context.remote_cas_url = remote
+    except InvalidArgumentError as e:
+        click.echo("Error: {}.".format(e), err=True)
+        sys.exit(-1)
 
-        context.cas_channel = context.channel
-
-        context.cas_client_key = context.client_key
-        context.cas_client_cert = context.client_cert
-        context.cas_server_cert = context.server_cert
+    context.update_period = update_period
+    context.parent = parent
 
     bot_interface = interface.BotInterface(context.channel)