datasette.create_token() method, closes #1951

datasette/app.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Sequence, Union, Tuple, Optional
+from typing import Sequence, Union, Tuple, Optional, Dict, Iterable
 import asgi_csrf
 import collections
 import datetime
@@ -16,6 +16,7 @@ import re
 import secrets
 import sys
 import threading
+import time
 import urllib.parse
 from concurrent import futures
 from pathlib import Path
@@ -465,6 +466,45 @@ class Datasette:
     def unsign(self, signed, namespace="default"):
         return URLSafeSerializer(self._secret, namespace).loads(signed)
 
+    def create_token(
+        self,
+        actor_id: str,
+        *,
+        expires_after: Optional[int] = None,
+        restrict_all: Optional[Iterable[str]] = None,
+        restrict_database: Optional[Dict[str, Iterable[str]]] = None,
+        restrict_resource: Optional[Dict[str, Dict[str, Iterable[str]]]] = None,
+    ):
+        token = {"a": actor_id, "token": "dstok", "t": int(time.time())}
+        if expires_after:
+            token["d"] = expires_after
+
+        def abbreviate_action(action):
+            # rename to abbr if possible
+            permission = self.permissions.get(action)
+            if not permission:
+                return action
+            return permission.abbr or action
+
+        if expires_after:
+            token["d"] = expires_after
+        if restrict_all or restrict_database or restrict_resource:
+            token["_r"] = {}
+            if restrict_all:
+                token["_r"]["a"] = [abbreviate_action(a) for a in restrict_all]
+            if restrict_database:
+                token["_r"]["d"] = {}
+                for database, actions in restrict_database.items():
+                    token["_r"]["d"][database] = [abbreviate_action(a) for a in actions]
+            if restrict_resource:
+                token["_r"]["r"] = {}
+                for database, resources in restrict_resource.items():
+                    for resource, actions in resources.items():
+                        token["_r"]["r"].setdefault(database, {})[resource] = [
+                            abbreviate_action(a) for a in actions
+                        ]
+        return "dstok_{}".format(self.sign(token, namespace="token"))
+
     def get_database(self, name=None, route=None):
         if route is not None:
             matches = [db for db in self.databases.values() if db.route == route]

datasette/cli.py

@@ -11,7 +11,6 @@ from runpy import run_module
 import shutil
 from subprocess import call
 import sys
-import time
 import webbrowser
 from .app import (
     OBSOLETE_SETTINGS,
@@ -730,41 +729,40 @@ def create_token(
     loop = asyncio.get_event_loop()
     loop.run_until_complete(ds.invoke_startup())
 
-    def fix_action(action):
-        # Warn if invalid, rename to abbr if possible
-        permission = ds.permissions.get(action)
-        if not permission:
-            # Output red message
+    # Warn about any unknown actions
+    actions = []
+    actions.extend(alls)
+    actions.extend([p[1] for p in databases])
+    actions.extend([p[2] for p in resources])
+    for action in actions:
+        if not ds.permissions.get(action):
             click.secho(
-                f" Unknown permission: {action} ",
+                f"  Unknown permission: {action} ",
                 fg="red",
                 err=True,
             )
-            return action
-        return permission.abbr or action
 
-    bits = {"a": id, "token": "dstok", "t": int(time.time())}
-    if expires_after:
-        bits["d"] = expires_after
-    if alls or databases or resources:
-        bits["_r"] = {}
-        if alls:
-            bits["_r"]["a"] = [fix_action(a) for a in alls]
-        if databases:
-            bits["_r"]["d"] = {}
-            for database, action in databases:
-                bits["_r"]["d"].setdefault(database, []).append(fix_action(action))
-        if resources:
-            bits["_r"]["r"] = {}
-            for database, table, action in resources:
-                bits["_r"]["r"].setdefault(database, {}).setdefault(table, []).append(
-                    fix_action(action)
-                )
-    token = ds.sign(bits, namespace="token")
-    click.echo("dstok_{}".format(token))
+    restrict_database = {}
+    for database, action in databases:
+        restrict_database.setdefault(database, []).append(action)
+    restrict_resource = {}
+    for database, resource, action in resources:
+        restrict_resource.setdefault(database, {}).setdefault(resource, []).append(
+            action
+        )
+
+    token = ds.create_token(
+        id,
+        expires_after=expires_after,
+        restrict_all=alls,
+        restrict_database=restrict_database,
+        restrict_resource=restrict_resource,
+    )
+    click.echo(token)
     if debug:
+        encoded = token[len("dstok_") :]
         click.echo("\nDecoded:\n")
-        click.echo(json.dumps(ds.unsign(token, namespace="token"), indent=2))
+        click.echo(json.dumps(ds.unsign(encoded, namespace="token"), indent=2))
 
 
 pm.hook.register_commands(cli=cli)