diff --git a/README.md b/README.md index baf7308..4c23541 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,30 @@ Options: ``` +## Authentication + +`dclient` can handle API tokens for Datasette instances that require authentication. + +You can pass an API token to `query` using `-t/--token` like this: + +```bash +dclient query https://latest.datasette.io/fixtures "select * from facetable" -t dstok_mytoken +``` +You can also store tokens for a specific URL prefix. To always use `dstok_mytoken` for any URL on the `https://latest.datasette.io/` instance you can run this: +```bash +dclient auth add https://latest.datasette.io/ +``` +Then paste in the token and hit enter when prompted to do so. + +To list which URLs you have set tokens for, run the `auth list` command: +```bash +dclient auth list +``` +To delete the token for a specific URL, run `auth remove`: +```bash +dclient auth remove https://latest.datasette.io/ +``` + ## Aliases You can assign an alias to a Datasette database using the `dclient alias` command: diff --git a/dclient/cli.py b/dclient/cli.py index db5b510..ac660f4 100644 --- a/dclient/cli.py +++ b/dclient/cli.py @@ -2,6 +2,7 @@ import click import httpx import json import pathlib +from .utils import token_for_url def get_config_dir(): @@ -17,7 +18,8 @@ def cli(): @cli.command() @click.argument("url") @click.argument("sql") -def query(url, sql): +@click.option("--token", "-t", help="API token") +def query(url, sql, token): """ Run a SQL query against a Datasette database URL @@ -29,17 +31,50 @@ def query(url, sql): url = aliases[url] if not url.endswith(".json"): url += ".json" - response = httpx.get(url, params={"sql": sql, "_shape": "objects"}) - if response.status_code != 200 or not response.json()["ok"]: - data = response.json() + if token is None: + # Maybe there's a token in auth.json? + token = token_for_url(url, _load_auths(get_config_dir() / "auth.json")) + headers = {} + if token: + headers["Authorization"] = f"Bearer {token}" + response = httpx.get(url, params={"sql": sql, "_shape": "objects"}, headers=headers) + if response.status_code != 200: + # Is it valid JSON? + try: + data = response.json() + except json.JSONDecodeError: + raise click.ClickException( + "{} status code. Response was not valid JSON".format( + response.status_code + ) + ) bits = [] if data.get("title"): bits.append(data["title"]) if data.get("error"): bits.append(data["error"]) + raise click.ClickException( + "{} status code. {}".format(response.status_code, ": ".join(bits)) + ) + + # We should have JSON now + try: + data = response.json() + except json.JSONDecodeError: + raise click.ClickException("Response was not valid JSON") + # ... but it may have a {"ok": false} error + if not data.get("ok"): + bits = [] + if data.get("title"): + bits.append(data["title"]) + if data.get("error"): + bits.append(data["error"]) + if not bits: + bits = [json.dumps(data)] raise click.ClickException(": ".join(bits)) - else: - click.echo(json.dumps(response.json()["rows"], indent=2)) + + # Output results + click.echo(json.dumps(response.json()["rows"], indent=2)) @cli.group() @@ -60,10 +95,10 @@ def list_(_json): click.echo(f"{alias} = {url}") -@alias.command() +@alias.command(name="add") @click.argument("name") @click.argument("url") -def add(name, url): +def alias_add(name, url): "Add an alias" config_dir = get_config_dir() config_dir.mkdir(parents=True, exist_ok=True) @@ -73,9 +108,9 @@ def add(name, url): aliases_file.write_text(json.dumps(aliases, indent=4)) -@alias.command() +@alias.command(name="remove") @click.argument("name") -def remove(name): +def alias_remove(name): "Remove an alias" config_dir = get_config_dir() aliases_file = config_dir / "aliases.json" @@ -87,9 +122,75 @@ def remove(name): raise click.ClickException("No such alias") +@cli.group() +def auth(): + "Manage authentication for different instances" + + +@auth.command(name="add") +@click.argument("alias_or_url") +@click.option("--token", prompt=True, hide_input=True) +def auth_add(alias_or_url, token): + """ + Add an authentication token for an alias or URL + + Example usage: + + \b + dclient auth add https://datasette.io/content + + Paste in the token when prompted. + """ + aliases_file = get_config_dir() / "aliases.json" + aliases = _load_aliases(aliases_file) + url = alias_or_url + if alias_or_url in aliases: + url = aliases[alias_or_url] + config_dir = get_config_dir() + config_dir.mkdir(parents=True, exist_ok=True) + auth_file = config_dir / "auth.json" + auths = _load_auths(auth_file) + auths[url] = token + auth_file.write_text(json.dumps(auths, indent=4)) + + +@auth.command(name="list") +def auth_list(): + "List stored API tokens" + auths_file = get_config_dir() / "auth.json" + click.echo("Tokens file: {}".format(auths_file)) + auths = _load_auths(auths_file) + if auths: + click.echo() + for url, token in auths.items(): + click.echo("{}:\t{}..".format(url, token[:1])) + + +@auth.command(name="remove") +@click.argument("alias_or_url") +def auth_remove(alias_or_url): + "Remove the API token for an alias or URL" + config_dir = get_config_dir() + auth_file = config_dir / "auth.json" + auths = _load_auths(auth_file) + try: + del auths[alias_or_url] + auth_file.write_text(json.dumps(auths, indent=4)) + except KeyError: + raise click.ClickException("No such URL or alias") + + def _load_aliases(aliases_file): if aliases_file.exists(): aliases = json.loads(aliases_file.read_text()) else: aliases = {} return aliases + + +def _load_auths(auth_file): + if auth_file.exists(): + auths = json.loads(auth_file.read_text()) + else: + auths = {} + return auths diff --git a/dclient/utils.py b/dclient/utils.py new file mode 100644 index 0000000..0a6576d --- /dev/null +++ b/dclient/utils.py @@ -0,0 +1,42 @@ +from urllib.parse import urlparse + + +def url_matches_prefix(url: str, prefix: str) -> bool: + url_parsed = urlparse(url) + prefix_parsed = urlparse(prefix) + + if ( + url_parsed.scheme != prefix_parsed.scheme + or url_parsed.netloc != prefix_parsed.netloc + ): + return False + + if url_parsed.path == prefix_parsed.path: + return True + + url_path = url_parsed.path + + # Add '/' to the end of the paths if they don't already end with '/' + prefix_path = prefix_parsed.path + + # Special treatment for /foo.json + if url_path.endswith(".json"): + if prefix_path + ".json" == url_path: + return True + + if not prefix_path.endswith("/"): + prefix_path += "/" + + return url_path.startswith(prefix_path) + + +def token_for_url(url: str, tokens: dict) -> str: + matches = [] + for prefix_url, token in tokens.items(): + if url_matches_prefix(url, prefix_url): + matches.append((prefix_url, token)) + # Sort by length of prefix_url, so that the longest match is first + matches.sort(key=lambda x: len(x[0]), reverse=True) + if matches: + return matches[0][1] + return None diff --git a/tests/test_cli_auth.py b/tests/test_cli_auth.py new file mode 100644 index 0000000..ba30e22 --- /dev/null +++ b/tests/test_cli_auth.py @@ -0,0 +1,41 @@ +from click.testing import CliRunner +from dclient.cli import cli +import pathlib +import json + + +def test_auth(mocker, tmpdir): + mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir)) + runner = CliRunner() + result = runner.invoke(cli, ["auth", "list"]) + assert result.exit_code == 0 + assert result.output.startswith("Tokens file:") + # Should only have one line + assert len([line for line in result.output.split("\n") if line.strip()]) == 1 + + # Now add a token + result2 = runner.invoke(cli, ["auth", "add", "https://example.com"], input="xyz\n") + assert result2.exit_code == 0 + + # Check the tokens file + auth_file = pathlib.Path(tmpdir) / "auth.json" + assert json.loads(auth_file.read_text()) == {"https://example.com": "xyz"} + + # auth list should show that now + result3 = runner.invoke(cli, ["auth", "list"]) + assert result3.output.startswith("Tokens file:") + assert "https://example.com" in result3.output + + # Remove should fail with an incorrect URL + result4 = runner.invoke(cli, ["auth", "remove", "https://example.com/foo"]) + assert result4.exit_code == 1 + assert result4.output == "Error: No such URL or alias\n" + + # Remove should work with the correct URL + result5 = runner.invoke(cli, ["auth", "remove", "https://example.com"]) + assert result5.exit_code == 0 + assert result5.output == "" + + # Check the tokens file + auth_file = pathlib.Path(tmpdir) / "auth.json" + assert json.loads(auth_file.read_text()) == {} diff --git a/tests/test_query.py b/tests/test_query.py index d565750..552acb0 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -2,6 +2,7 @@ from click.testing import CliRunner from dclient.cli import cli import json import pathlib +import pytest def test_query_error(httpx_mock): @@ -17,10 +18,14 @@ def test_query_error(httpx_mock): runner = CliRunner() result = runner.invoke(cli, ["query", "https://example.com", "hello"]) assert result.exit_code == 1 - assert result.output == "Error: Invalid SQL: Statement must be a SELECT\n" + assert ( + result.output + == "Error: 400 status code. Invalid SQL: Statement must be a SELECT\n" + ) -def test_query(httpx_mock): +@pytest.mark.parametrize("with_token", (False, True)) +def test_query(httpx_mock, with_token): httpx_mock.add_response( json={ "ok": True, @@ -37,10 +42,22 @@ def test_query(httpx_mock): status_code=200, ) runner = CliRunner() - result = runner.invoke(cli, ["query", "https://example.com", "hello"]) + args = ["query", "https://example.com", "hello"] + if with_token: + args.append("--token") + args.append("xyz") + result = runner.invoke(cli, args) assert result.exit_code == 0 assert json.loads(result.output) == [{"5 * 2": 10}] + # Check the request + request = httpx_mock.get_request() + assert str(request.url) == "https://example.com.json?sql=hello&_shape=objects" + if with_token: + assert request.headers["authorization"] == "Bearer xyz" + else: + assert "authorization" not in request.headers + def test_aliases(mocker, tmpdir, httpx_mock): mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir)) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..81a63e8 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,67 @@ +from dclient.utils import token_for_url, url_matches_prefix +import pytest + + +@pytest.mark.parametrize( + "url,prefix_url,expected", + ( + ("https://example.com/foo/bar", "https://example.com/foo", True), + ("https://example.com/foo/bar2", "https://example.com/foo/bar", False), + ("https://example.com/foo/bar/baz", "https://example.com/foo/bar", True), + ("https://example.com/foo/bar/baz", "https://example.com/foo", True), + ("https://example.com/foo.json", "https://example.com/foo", True), + # different scheme + ( + "http://example.com/foo/bar", + "https://example.com/foo", + False, + ), + # different netloc + ( + "https://example.org/foo/bar", + "https://example.com/foo", + False, + ), + # exactly the same + ("https://example.com/foo", "https://example.com/foo", True), + # trailing '/' + ( + "https://example.com/foo/bar", + "https://example.com/foo/bar/", + False, + ), + ( + "https://example.com/foo/bar/baz", + "https://example.com/foo/bar/", + True, + ), + ), +) +def test_url_matches_prefix(url, prefix_url, expected): + assert url_matches_prefix(url, prefix_url) == expected + + +@pytest.mark.parametrize( + "url,tokens,expected", + ( + ("https://foo.com/bar", {"https://foo.com": "foo"}, "foo"), + ("https://foo.com/bar", {"https://foo.com/baz": "baz"}, None), + ( + "https://foo.com/bar", + {"https://foo.com": "foo", "https://foo.com/bar": "bar"}, + "bar", + ), + ( + "https://foo.com/bar/baz", + { + "https://foo.com": "foo", + "https://foo.com/bar/baz": "baz", + "https://foo.com/bar": "bar", + }, + "baz", + ), + ), +) +def test_token_for_url(url, tokens, expected): + # Should always return longest matching of the available options + assert token_for_url(url, tokens) == expected