def _run(flags):
"""Runs the main uploader program given parsed flags.
Args:
flags: An `argparse.Namespace`.
"""
logging.set_stderrthreshold(logging.WARNING)
intent = _get_intent(flags)
store = auth.CredentialsStore()
if isinstance(intent, _AuthRevokeIntent):
store.clear()
sys.stderr.write("Logged out of uploader.\n")
sys.stderr.flush()
return
# TODO(b/141723268): maybe reconfirm Google Account prior to reuse.
credentials = store.read_credentials()
if not credentials:
_prompt_for_user_ack(intent)
client_config = json.loads(auth.OAUTH_CLIENT_CONFIG)
flow = auth.build_installed_app_flow(client_config)
credentials = flow.run(force_console=flags.auth_force_console)
sys.stderr.write("\n") # Extra newline after auth flow messages.
store.write_credentials(credentials)
channel_options = None
if flags.grpc_creds_type == "local":
channel_creds = grpc.local_channel_credentials()
elif flags.grpc_creds_type == "ssl":
channel_creds = grpc.ssl_channel_credentials()
elif flags.grpc_creds_type == "ssl_dev":
channel_creds = grpc.ssl_channel_credentials(dev_creds.DEV_SSL_CERT)
channel_options = [("grpc.ssl_target_name_override", "localhost")]
else:
msg = "Invalid --grpc_creds_type %s" % flags.grpc_creds_type
raise base_plugin.FlagsError(msg)
try:
server_info = _get_server_info(flags)
except server_info_lib.CommunicationError as e:
_die(str(e))
_handle_server_info(server_info)
if not server_info.api_server.endpoint:
logging.error("Server info response: %s", server_info)
_die("Internal error: frontend did not specify an API server")
composite_channel_creds = grpc.composite_channel_credentials(
channel_creds, auth.id_token_call_credentials(credentials)
)
# TODO(@nfelt): In the `_UploadIntent` case, consider waiting until
# logdir exists to open channel.
channel = grpc.secure_channel(
server_info.api_server.endpoint,
composite_channel_creds,
options=channel_options,
)
with channel:
intent.execute(server_info, channel)
def test_clear_nonexistent_file(self):
root = self.get_temp_dir()
path = os.path.join(root, "tensorboard", "credentials",
"uploader-creds.json")
self.assertFalse(os.path.exists(path))
auth.CredentialsStore(user_config_directory=root).clear()
self.assertFalse(os.path.exists(path))
def __init__(self,
logdir,
name=None,
description=None,
origin="",
api_endpoint=""):
self.logdir = logdir
self.name = name
self.description = description
self.origin = name
self.api_endpoint = description
self.args = Namespace(logdir=self.logdir,
name=self.name,
description=self.description,
origin=self.origin,
api_endpoint=self.api_endpoint)
store = auth.CredentialsStore()
credentials = store.read_credentials()
composite_channel_creds = grpc.composite_channel_credentials(
grpc.ssl_channel_credentials(),
auth.id_token_call_credentials(credentials))
self.server_info = _get_server_info(self.args)
self.channel = grpc.secure_channel(
self.server_info.api_server.endpoint,
composite_channel_creds,
options=None)
def test_clear_existent_file(self):
root = self.get_temp_dir()
path = os.path.join(root, "tensorboard", "credentials",
"uploader-creds.json")
os.makedirs(os.path.dirname(path))
open(path, mode="w").close()
self.assertTrue(os.path.exists(path))
auth.CredentialsStore(user_config_directory=root).clear()
self.assertFalse(os.path.exists(path))
def test_read_invalid_json_file(self):
root = self.get_temp_dir()
store = auth.CredentialsStore(user_config_directory=root)
path = os.path.join(root, "tensorboard", "credentials",
"uploader-creds.json")
os.makedirs(os.path.dirname(path))
with open(path, mode="w") as f:
f.write("{}")
with self.assertRaises(ValueError):
store.read_credentials()
def get_api_client(api_endpoint=None):
server_info = _get_server_info(api_endpoint=api_endpoint)
_handle_server_info(server_info)
channel_creds = grpc.ssl_channel_credentials()
credentials = auth.CredentialsStore().read_credentials()
if credentials:
channel_creds = grpc.composite_channel_credentials(
channel_creds, auth.id_token_call_credentials(credentials))
channel = grpc.secure_channel(server_info.api_server.endpoint,
channel_creds)
return export_service_pb2_grpc.TensorBoardExporterServiceStub(channel)
def _run(flags, experiment_url_callback=None):
"""Runs the main uploader program given parsed flags.
Args:
flags: An `argparse.Namespace`.
experiment_url_callback: A function accepting a single string argument
containing the full TB.dev URL of the uploaded experiment.
"""
logging.set_stderrthreshold(logging.WARNING)
intent = _get_intent(flags, experiment_url_callback)
store = auth.CredentialsStore()
if isinstance(intent, _AuthRevokeIntent):
store.clear()
sys.stderr.write("Logged out of uploader.\n")
sys.stderr.flush()
return
# TODO(b/141723268): maybe reconfirm Google Account prior to reuse.
credentials = store.read_credentials()
if not credentials:
_prompt_for_user_ack(intent)
client_config = json.loads(auth.OAUTH_CLIENT_CONFIG)
flow = auth.build_installed_app_flow(client_config)
credentials = flow.run(force_console=flags.auth_force_console)
sys.stderr.write("\n") # Extra newline after auth flow messages.
store.write_credentials(credentials)
(channel_creds, channel_options) = flags.grpc_creds_type.channel_config()
try:
server_info = _get_server_info(flags)
except server_info_lib.CommunicationError as e:
_die(str(e))
_handle_server_info(server_info)
logging.info("Received server info: <%r>", server_info)
if not server_info.api_server.endpoint:
logging.error("Server info response: %s", server_info)
_die("Internal error: frontend did not specify an API server")
composite_channel_creds = grpc.composite_channel_credentials(
channel_creds, auth.id_token_call_credentials(credentials))
# TODO(@nfelt): In the `_UploadIntent` case, consider waiting until
# logdir exists to open channel.
channel = grpc.secure_channel(
server_info.api_server.endpoint,
composite_channel_creds,
options=channel_options,
)
with channel:
intent.execute(server_info, channel)
def test_write_creates_private_file(self):
root = self.get_temp_dir()
auth.CredentialsStore(user_config_directory=root).write_credentials(
google.oauth2.credentials.Credentials(token=None,
refresh_token="12345"))
path = os.path.join(root, "tensorboard", "credentials",
"uploader-creds.json")
self.assertTrue(os.path.exists(path))
# Skip permissions check on Windows.
if os.name != "nt":
self.assertEqual(0o600, os.stat(path).st_mode & 0o777)
with open(path) as f:
contents = json.load(f)
self.assertEqual("12345", contents["refresh_token"])
def _run(flags):
"""Runs the main uploader program given parsed flags.
Args:
flags: An `argparse.Namespace`.
"""
logging.set_stderrthreshold(logging.WARNING)
intent = _get_intent(flags)
store = auth.CredentialsStore()
if isinstance(intent, _AuthRevokeIntent):
store.clear()
sys.stderr.write('Logged out of uploader.\n')
sys.stderr.flush()
return
# TODO(b/141723268): maybe reconfirm Google Account prior to reuse.
credentials = store.read_credentials()
if not credentials:
_prompt_for_user_ack(intent)
client_config = json.loads(auth.OAUTH_CLIENT_CONFIG)
flow = auth.build_installed_app_flow(client_config)
credentials = flow.run(force_console=flags.auth_force_console)
sys.stderr.write('\n') # Extra newline after auth flow messages.
store.write_credentials(credentials)
channel_options = None
if flags.grpc_creds_type == 'local':
channel_creds = grpc.local_channel_credentials()
elif flags.grpc_creds_type == 'ssl':
channel_creds = grpc.ssl_channel_credentials()
elif flags.grpc_creds_type == 'ssl_dev':
channel_creds = grpc.ssl_channel_credentials(dev_creds.DEV_SSL_CERT)
channel_options = [('grpc.ssl_target_name_override', 'localhost')]
else:
msg = 'Invalid --grpc_creds_type %s' % flags.grpc_creds_type
raise base_plugin.FlagsError(msg)
composite_channel_creds = grpc.composite_channel_credentials(
channel_creds, auth.id_token_call_credentials(credentials))
# TODO(@nfelt): In the `_UploadIntent` case, consider waiting until
# logdir exists to open channel.
channel = grpc.secure_channel(flags.endpoint,
composite_channel_creds,
options=channel_options)
with channel:
intent.execute(channel)
def test_write_and_read_roundtrip(self):
orig_creds = google.oauth2.credentials.Credentials(
token="12345",
refresh_token="67890",
token_uri="https://oauth2.googleapis.com/token",
client_id="my-client",
client_secret="123abc456xyz",
scopes=["userinfo", "email"])
root = self.get_temp_dir()
store = auth.CredentialsStore(user_config_directory=root)
store.write_credentials(orig_creds)
creds = store.read_credentials()
self.assertEqual(orig_creds.refresh_token, creds.refresh_token)
self.assertEqual(orig_creds.token_uri, creds.token_uri)
self.assertEqual(orig_creds.client_id, creds.client_id)
self.assertEqual(orig_creds.client_secret, creds.client_secret)
def test_write_overwrites_file(self):
root = self.get_temp_dir()
store = auth.CredentialsStore(user_config_directory=root)
# Write twice to ensure that we're overwriting correctly.
store.write_credentials(
google.oauth2.credentials.Credentials(token=None,
refresh_token="12345"))
store.write_credentials(
google.oauth2.credentials.Credentials(token=None,
refresh_token="67890"))
path = os.path.join(root, "tensorboard", "credentials",
"uploader-creds.json")
self.assertTrue(os.path.exists(path))
with open(path) as f:
contents = json.load(f)
self.assertEqual("67890", contents["refresh_token"])
def test_write_wrong_type(self):
creds = google.auth.credentials.AnonymousCredentials()
with self.assertRaisesRegex(TypeError, "google.auth.credentials"):
auth.CredentialsStore(user_config_directory=None).write_credentials(
creds
)
def test_no_config_dir(self):
store = auth.CredentialsStore(user_config_directory=None)
self.assertIsNone(store.read_credentials())
creds = google.oauth2.credentials.Credentials(token=None)
store.write_credentials(creds)
store.clear()
def test_read_nonexistent_file(self):
root = self.get_temp_dir()
store = auth.CredentialsStore(user_config_directory=root)
self.assertIsNone(store.read_credentials())