diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index fae1480f1588..8e26c94d4009 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -152,26 +152,58 @@ def make_url(self, endpoint, path, params): location = endpoint.host_with_port return urllib.parse.urlunsplit((endpoint.protocol, location, path, urllib.parse.urlencode(params), '')) - def parse_token(self, token_file): + def parse_token(self, token_file, iam_token_file=None): if token_file: - self.token = token_file.readline().rstrip('\r\n') + self.token_type, self.token = self.read_token_from_file(token_file, 'OAuth') token_file.close() - if self.token is None: - self.token = os.getenv('YDB_TOKEN') + return + + if iam_token_file: + self.token_type, self.token = self.read_token_from_file(iam_token_file, 'Bearer') + iam_token_file.close() + return + + token_value = os.getenv('YDB_TOKEN') + if token_value is not None: + self.token_type, self.token = self.parse_token_value(token_value, 'OAuth') + return + + token_value = os.getenv('IAM_TOKEN') + if token_value is not None: + self.token_type, self.token = self.parse_token_value(token_value, 'Bearer') + return + + default_token_paths = [ + ('OAuth', os.path.expanduser(os.path.join('~', '.ydb', 'token'))), + ('Bearer', os.path.expanduser(os.path.join('~', '.ydb', 'iam_token'))), + ] + for token_type, token_file_path in default_token_paths: + self.token_type, self.token = self.read_token_file(token_file_path, token_type) if self.token is not None: - self.token = self.token.strip() - if self.token is None: - try: - path = os.path.expanduser(os.path.join('~', '.ydb', 'token')) - with open(path) as f: - self.token = f.readline().strip('\r\n') - except Exception: - pass - - if self.token is not None and len(self.token.split(' ')) == 2: - self.token_type, self.token = self.token.split(' ') + return + + def read_token_from_file(self, token_file, default_token_type): + if token_file is None: + return default_token_type, None + token_value = token_file.readline().rstrip('\r\n') + return self.parse_token_value(token_value, default_token_type) + + def read_token_file(self, token_file_path, default_token_type): + if token_file_path is None: + return default_token_type, None + try: + return self.read_token_from_file_and_close(open(token_file_path, 'r'), default_token_type) + except Exception: + return default_token_type, None + + def parse_token_value(self, token_value, default_token_type): + if token_value is None: + return default_token_type, None + splitted = token_value.strip().split(' ') + if len(splitted) == 2: + return splitted else: - self.token_type = 'OAuth' + return default_token_type, token_value def apply_args(self, args, with_localhost=True): self.args = args @@ -199,7 +231,7 @@ def apply_args(self, args, with_localhost=True): if 'http' not in protocols and 'https' in protocols: self.mon_protocol = 'https' - self.parse_token(args.token_file) + self.parse_token(args.token_file, args.iam_token_file) self.domain = 1 self.verbose = args.verbose or args.debug self.debug = args.debug @@ -218,7 +250,9 @@ def add_host_access_options(self, parser, with_endpoint=True): g.add_argument('--endpoint', '-e', metavar='[PROTOCOL://]HOST[:PORT]', type=str, required=True, action='append', help=ConnectionParams.ENDPOINT_HELP) g.add_argument('--grpc-port', type=int, default=2135, metavar='PORT', help='GRPC port to use for procedure invocation') g.add_argument('--mon-port', type=int, default=8765, metavar='PORT', help='HTTP monitoring port for viewer JSON access') - g.add_argument('--token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to token file') + token_group = g.add_mutually_exclusive_group() + token_group.add_argument('--token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to token file') + token_group.add_argument('--iam-token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to IAM token file') g.add_argument('--ca-file', metavar='PATH', dest='cafile', type=str, help='File containing PEM encoded root certificates for SSL/TLS connections. ' 'If this parameter is empty, the default roots will be used.') g.add_argument('--http-timeout', type=int, default=5, help='Timeout for blocking socket I/O operations during HTTP(s) queries')