stash

cli-client/toolshed-client.py

@@ -1,9 +1,17 @@
 #!/usr/bin/env python3
 import argparse
+import base64
 import os
+import sys
 import requests
 from nacl.signing import SigningKey
-from json import dumps
+from json import dumps, loads
+
+COMMANDS = {
+    'getinventory': {'path': '/api/inventory_items/{handle}/', 'method': 'get'},
+    'additem': {'path': '/api/inventory_items/{handle}/', 'method': 'post'},
+    'delitem': {'path': '/api/inventory_items/{handle}/{internal_id}/', 'method': 'delete'},
+}
 
 
 class ToolshedApi:
@@ -29,6 +37,14 @@ class ToolshedApi:
         self.user = user
         self.host = host
         self.signing_key = signing_key
+        self._spec = None
+
+    def get_spec(self):
+        if self._spec is None:
+            response = requests.get("http://" + self.host + "/docs/?format=openapi")
+            response.raise_for_status()
+            self._spec = response.json()
+        return self._spec
 
     def get(self, target):
         url = "http://" + self.host + target
@@ -45,6 +61,175 @@ class ToolshedApi:
         response = requests.post(url, headers={"Authorization": "Signature " + self.user + ":" + signature}, json=data)
         return response.json()
 
+    def delete(self, target):
+        url = "http://" + self.host + target
+        signed = self.signing_key.sign(url.encode('utf-8'))
+        signature = signed.signature.hex()
+        response = requests.delete(url, headers={"Authorization": "Signature " + self.user + ":" + signature})
+        if not response.content:
+            return {"deleted": response.ok}
+        return response.json()
+
+    def get_raw(self, target):
+        """Like get(), but returns the raw response body instead of parsing it as JSON - for
+        endpoints like /api/export/ that hand back a zip file, not a JSON document."""
+        url = "http://" + self.host + target
+        signed = self.signing_key.sign(url.encode('utf-8'))
+        signature = signed.signature.hex()
+        response = requests.get(url, headers={"Authorization": "Signature " + self.user + ":" + signature})
+        response.raise_for_status()
+        return response.content
+
+
+def resolve_schema(spec, schema):
+    if '$ref' in schema:
+        return spec['definitions'][schema['$ref'].split('/')[-1]]
+    return schema
+
+
+def operation_parameters(spec, path, method):
+    """Path- and body-parameters for an operation, per the endpoint's own OpenAPI spec entry."""
+    path_item = spec['paths'][path]
+    params = path_item.get('parameters', []) + path_item[method].get('parameters', [])
+    path_params = [p for p in params if p.get('in') == 'path']
+    body_param = next((p for p in params if p.get('in') == 'body'), None)
+    return path_params, body_param
+
+
+def parse_kv_args(cmd_args):
+    """['name=Drill', 'owned_quantity=2'] -> {'name': 'Drill', 'owned_quantity': '2'}"""
+    pairs = {}
+    for arg in cmd_args:
+        if '=' in arg:
+            key, value = arg.split('=', 1)
+            pairs[key] = value
+    return pairs
+
+
+def resolve_path_params(api, path_params, cmd_args, json_input):
+    values = {}
+    positional = [arg for arg in cmd_args if '=' not in arg]
+    kv = parse_kv_args(cmd_args)
+    for param in path_params:
+        name = param['name']
+        if name == 'handle':
+            values[name] = api.user
+            continue
+        if json_input is not None:
+            value = json_input.get(name, json_input.get('id'))
+            if value is None:
+                raise ValueError("Missing required path parameter '{}' in --json stdin input".format(name))
+        elif name in kv:
+            value = kv[name]
+        elif 'id' in kv:
+            value = kv['id']
+        elif positional:
+            value = positional.pop(0)
+        else:
+            value = input("{}: ".format(name))
+        values[name] = value
+    return values
+
+
+def resolve_body(spec, body_param, cmd_args, json_input):
+    if body_param is None:
+        return None
+    if json_input is not None:
+        return json_input
+    kv = parse_kv_args(cmd_args)
+    if kv:
+        return kv
+    schema = resolve_schema(spec, body_param['schema'])
+    writable_fields = [name for name, prop in schema.get('properties', {}).items() if not prop.get('readOnly')]
+    body = {}
+    for field in writable_fields:
+        value = input("{}: ".format(field))
+        if value != '':
+            body[field] = value
+    return body
+
+
+def build_request(api, spec, command, cmd_args, json_input):
+    path_params, body_param = operation_parameters(spec, command['path'], command['method'])
+    values = resolve_path_params(api, path_params, cmd_args, json_input)
+    url = command['path']
+    for name, value in values.items():
+        url = url.replace('{' + name + '}', str(value))
+    body = resolve_body(spec, body_param, cmd_args, json_input)
+    return url, body
+
+
+def run_command(api, cmd, cmd_args, json_input):
+    if cmd == 'export':
+        path = cmd_args[0] if cmd_args else 'toolshed-export.zip'
+        data = api.get_raw("/api/export/")
+        with open(path, 'wb') as f:
+            f.write(data)
+        return {'exported_to': path, 'bytes': len(data)}
+    elif cmd == 'import':
+        if not cmd_args:
+            raise ValueError("import requires the path to a previously exported zip file")
+        path = cmd_args[0]
+        with open(path, 'rb') as f:
+            data = f.read()
+        return api.post("/api/import/", {"zip": base64.b64encode(data).decode('ascii')})
+
+    command = COMMANDS.get(cmd)
+    if command is None:
+        raise ValueError("Unknown command: " + cmd)
+    spec = api.get_spec()
+    url, body = build_request(api, spec, command, cmd_args, json_input)
+    if command['method'] == 'get':
+        return api.get(url)
+    elif command['method'] == 'post':
+        return api.post(url, body or {})
+    elif command['method'] == 'delete':
+        return api.delete(url)
+    else:
+        raise ValueError("Unsupported method: " + command['method'])
+
+
+def read_json_input():
+    if sys.stdin.isatty():
+        return {}
+    raw = sys.stdin.read()
+    return loads(raw) if raw.strip() else {}
+
+
+def stringify_cell(value):
+    if value is None:
+        return ''
+    if isinstance(value, (list, tuple)):
+        return ', '.join(stringify_cell(v) for v in value)
+    if isinstance(value, dict):
+        return dumps(value)
+    return str(value)
+
+
+def format_table(rows):
+    if not rows:
+        return '(empty)'
+    columns = []
+    for row in rows:
+        for key in row.keys():
+            if key not in columns:
+                columns.append(key)
+    lines = ['| ' + ' | '.join(columns) + ' |', '| ' + ' | '.join('---' for _ in columns) + ' |']
+    for row in rows:
+        lines.append('| ' + ' | '.join(stringify_cell(row.get(column)) for column in columns) + ' |')
+    return '\n'.join(lines)
+
+
+def print_result(result, as_json):
+    if as_json:
+        print(dumps(result, indent=2))
+    elif isinstance(result, list) and all(isinstance(item, dict) for item in result):
+        print(format_table(result))
+    elif isinstance(result, dict):
+        print(format_table([result]))
+    else:
+        print(result)
+
 
 def main():
     host = os.environ.get('TOOLSHED_HOST')
@@ -55,7 +240,10 @@ def main():
     parser.add_argument('--host', help='Toolshed host')
     parser.add_argument('--user', help='Toolshed user')
     parser.add_argument('--key', help='Toolshed key')
+    parser.add_argument('--json', action='store_true',
+                        help='Read input as JSON from stdin and print output as JSON, instead of prompting/printing')
     parser.add_argument('cmd', help='Command')
+    parser.add_argument('args', nargs='*', help="Command arguments, as key=value pairs (e.g. name=Drill)")
     args = parser.parse_args()
 
     if args.host is not None:
@@ -69,14 +257,9 @@ def main():
 
     api = ToolshedApi(user, host, key)
 
-    if args.cmd == 'getinventory':
-        inv = api.get("/api/inventory_items/")
-        print(inv)
-    elif args.cmd == 'additem':
-        inv = api.post("/api/inventory_items/", {"name": "test"})
-        print(inv)
-    else:
-        print("Unknown command: " + args.cmd)
+    json_input = read_json_input() if args.json else None
+
+    print_result(run_command(api, args.cmd, args.args, json_input), args.json)
 
 
 if __name__ == '__main__':