feat: SSO Improvement - alter user_sessions table to include access token, implement CRUD ops, GET, POST, PATCH APIs and det token CLIs

harness/determined/cli/user.py

@@ -1,14 +1,26 @@
 import argparse
 import collections
+import functools
 import getpass
-from typing import Any, List
+import json
+from typing import Any, List, Sequence
 
 from determined import cli
 from determined.cli import errors, render
-from determined.common import api
+from determined.common import api, util
 from determined.common.api import authentication, bindings
 from determined.experimental import client
 
+TOKEN_HEADERS = [
+    "ID",
+    "User ID",
+    "Description",
+    "Created At",
+    "Expires At",
+    "Revoked",
+    "Token Type",
+]
+
 FullUser = collections.namedtuple(
     "FullUser",
     [
@@ -71,6 +83,19 @@ def log_in_user(args: argparse.Namespace) -> None:
     else:
         username = args.username
 
+    if args.token is not None:
+        token_store = authentication.TokenStore(args.master)
+        token_store.set_token(username, args.token)
+        token_store.set_active(username)
+
+        d = client.Determined._from_session(
+            api.Session(master=args.master, username=username, token=args.token, cert=cli.cert)
+        )
+        user = d.whoami()
+        if user.username != username:
+            raise errors.CliError("Token does not match the provided username")
+        return
+
     message = "Password for user '{}': ".format(username)
     password = getpass.getpass(message)
 
@@ -218,6 +243,110 @@ def edit(args: argparse.Namespace) -> None:
         raise errors.CliError("No field provided. Use 'det user edit -h' for usage.")
 
 
+def render_token_info(token_info: Sequence[bindings.v1TokenInfo]) -> None:
+    values = []
+    for ti in token_info:
+        value = [
+            ti.id,
+            ti.userId,
+            ti.description,
+            ti.createdAt,
+            ti.expiry,
+            ti.revoked,
+            ti.tokenType,
+        ]
+        values.append(value)
+    render.tabulate_or_csv(TOKEN_HEADERS, values, False)
+
+
+def print_token_info(data: Sequence[bindings.v1TokenInfo], args: argparse.Namespace) -> None:
+    if data:
+        json_data = [elem.to_json() for elem in data]
+        if len(data) == 1:
+            json_data = json_data[0]
+    else:
+        json_data = {}
+    if args.yaml:
+        print(util.yaml_safe_dump(json_data, default_flow_style=False))
+    else:
+        render.print_json(json_data)
+
+
+def describe_token(args: argparse.Namespace) -> None:
+    sess = cli.setup_session(args)
+    result = []
+    for arg in args.username:
+        userID = bindings.get_GetUserByUsername(session=sess, username=arg).user.id
+        resp = bindings.get_GetAccessToken(session=sess, userId=userID)
+        result.append(resp.tokenInfo)
+    render_token_info(result)
+
+
+def list_tokens(args: argparse.Namespace) -> None:
+    sess = cli.setup_session(args)
+    resp = bindings.get_GetAllAccessTokens(sess, includeInactive=args.all)
+    if args.json or args.yaml:
+        print_token_info(resp.tokenInfo, args)
+        return
+    render_token_info(resp.tokenInfo)
+
+
+def revoke_token(args: argparse.Namespace) -> None:
+    sess = cli.setup_session(args)
+    try:
+        request = bindings.v1PatchAccessTokenRequest(
+            tokenId=args.token_id, description=None, setRevoked=True
+        )
+        resp = bindings.patch_PatchAccessToken(sess, body=request, tokenId=args.token_id)
+        print(json.dumps(resp.to_json(), indent=2))
+    except api.errors.NotFoundException:
+        raise errors.CliError("Token not found")
+    print("Successfully updated token with ID: {}".format(args.token_id))
+
+
+def create_token(args: argparse.Namespace) -> None:
+    sess = cli.setup_session(args)
+    current_user, token_user = sess.username, args.username
+    request, handler = None, None
+    if token_user is None or token_user == current_user:
+        user_id = bindings.get_GetUserByUsername(session=sess, username=current_user).user.id
+    else:
+        user_id = bindings.get_GetUserByUsername(session=sess, username=token_user).user.id
+    request = bindings.v1PostAccessTokenRequest(
+        lifespan=args.lifespan, userId=user_id, description=args.description
+    )
+    handler = functools.partial(bindings.post_PostAccessToken, userId=user_id)
+    resp = handler(session=sess, body=request).to_json()
+
+    outputString = None
+    if args.yaml:
+        outputString = util.yaml_safe_dump(resp, default_flow_style=False)
+    elif args.json:
+        outputString = json.dumps(resp, indent=2)
+    else:
+        outputString = resp["token"]
+
+    if args.file_path:
+        with open(args.file_path, "w") as file:
+            file.write(outputString)
+    else:
+        print(outputString)
+
+
+def update_token(args: argparse.Namespace) -> None:
+    sess = cli.setup_session(args)
+    try:
+        if args.description and args.token_id:
+            request = bindings.v1PatchAccessTokenRequest(
+                tokenId=args.token_id, description=args.description, setRevoked=False
+            )
+            resp = bindings.patch_PatchAccessToken(sess, body=request, tokenId=args.token_id)
+        print(json.dumps(resp.to_json(), indent=2))
+    except api.errors.NotFoundException:
+        raise errors.CliError("Token not found")
+    print("Successfully updated token with ID: {}".format(args.token_id))
+
+
 AGENT_USER_GROUP_ARGS = [
     cli.Arg("--agent-uid", type=int, help="UID on the agent to run tasks as"),
     cli.Arg("--agent-user", help="user on the agent to run tasks as"),
@@ -238,7 +367,8 @@ def edit(args: argparse.Namespace) -> None:
             ),
         ], is_default=True),
         cli.Cmd("login", log_in_user, "log in user", [
-            cli.Arg("username", nargs="?", default=None, help="name of user to log in as")
+            cli.Arg("username", nargs="?", default=None, help="name of user to log in as"),
+            cli.Arg("--token", default=None, help="token to use for authentication"),
         ]),
         cli.Cmd("rename", rename, "change username for user", [
             cli.Arg(
@@ -313,6 +443,46 @@ def edit(args: argparse.Namespace) -> None:
                 help="grant/remove user admin permissions",
             ),
         ]),
+        cli.Cmd("token", None, "manage access tokens", [
+            cli.Cmd("describe", describe_token, "describe token info", [
+                cli.Arg("username", nargs=argparse.ONE_OR_MORE, default=None,
+                        help="name of user to describe token"),
+                cli.Group(
+                    cli.output_format_args["json"],
+                    cli.output_format_args["yaml"],
+                ),
+            ]),
+            cli.Cmd("list ls", list_tokens, "list all active access tokens", [
+                cli.Arg("--all", "-a", action="store_true", default=None,
+                        help="list all access tokens, including revoked & expired tokens"),
+                cli.Group(
+                    cli.output_format_args["json"],
+                    cli.output_format_args["yaml"],
+                ),
+            ]),
+            cli.Cmd("revoke", revoke_token, "revoke token", [
+                cli.Arg("token_id", help="revoke given access token"),
+            ]),
+            cli.Cmd("create", create_token, "create token", [
+                cli.Arg("--username", type=str, help="name of user to create token"),
+                cli.Arg("--lifespan", type=str, help="give expiry lifespan"),
+                cli.Arg("--description", type=str, default=None, help="description of new token"),
+                cli.Arg("--file-path", type=str, help="write token to file"),
+                cli.Group(
+                    cli.output_format_args["json"],
+                    cli.output_format_args["yaml"],
+                ),
+            ]),
+            cli.Cmd("update", update_token, "update token info", [
+                cli.Arg("token_id", help="revoke given access token"),
+                cli.Arg("description", type=str, default=None,
+                        help="description of token"),
+                cli.Group(
+                    cli.output_format_args["json"],
+                    cli.output_format_args["yaml"],
+                ),
+            ]),
+        ]),
     ])
 ]  # type: List[Any]