Add argument --iam-token-file in ydb-dstool (#20303)

1 files changed, 52 insertions, 18 deletions

diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py
index 7564f0a9619..db151c96c46 100644
--- a/ydb/apps/dstool/lib/common.py
+++ b/ydb/apps/dstool/lib/common.py
@@ -154,26 +154,58 @@ class ConnectionParams:
             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
@@ -201,7 +233,7 @@ class ConnectionParams:
         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
@@ -220,7 +252,9 @@ class ConnectionParams:
             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')