added queries subcommand which imports and exports custom queries

closes exploide#3

Author: exploide

src/bloodhound_cli/api/__init__.py

@@ -178,6 +178,26 @@ def add_to_asset_group(self, asset_group_id, sids):
         return self._send("POST", endpoint, data)
 
 
+    def get_saved_queries(self, sort_by=None):
+        """Get saved custom queries."""
+
+        endpoint = "/api/v2/saved-queries"
+        if sort_by is not None:
+            endpoint += f"?sort_by={urllib.parse.quote_plus(sort_by)}"
+        return self._send("GET", endpoint)
+
+
+    def add_saved_query(self, name, query):
+        """Add a custom query."""
+
+        endpoint = "/api/v2/saved-queries"
+        data = {
+            "name": name,
+            "query": query,
+        }
+        return self._send("POST", endpoint, data)
+
+
     def domains(self, collected=None):
         """Return available domains."""
 

src/bloodhound_cli/cli/__init__.py

@@ -9,6 +9,7 @@
 from .domains import domains
 from .groups import groups
 from .owned import owned
+from .queries import queries
 from .stats import stats
 from .upload import upload
 from .users import users
@@ -30,6 +31,7 @@ def bloodhound_cli(debug=False):
 bloodhound_cli.add_command(domains)
 bloodhound_cli.add_command(groups)
 bloodhound_cli.add_command(owned)
+bloodhound_cli.add_command(queries)
 bloodhound_cli.add_command(stats)
 bloodhound_cli.add_command(upload)
 bloodhound_cli.add_command(users)

src/bloodhound_cli/cli/queries.py

@@ -0,0 +1,62 @@
+import json
+import sys
+
+import click
+
+from bloodhound_cli.api.exceptions import ApiException
+from bloodhound_cli.api.from_config import api
+from bloodhound_cli.logger import log
+
+
+@click.command()
+@click.argument("file", type=click.Path(dir_okay=False))
+@click.option("--save", is_flag=True, help="Save existing custom queries to file.")
+def queries(file, save):
+    """Import and export custom queries.
+
+    Imports custom Cypher queries from a file to the BloodHound database.
+    If --save is specified, existing queries are exported to a file instead of imported.
+
+    For import, the file must either be in the format that --save produces or in the legacy Bloodhound's customqueries.json format.
+    However, not everything from the latter might be compatible.
+    """
+
+    if save:
+    # export queries
+        result = api.get_saved_queries(sort_by="name")
+        queries_to_export = [{key: entry[key] for key in ("name", "query")} for entry in result]
+        with open(file, "w", encoding="UTF-8") as f:
+            json.dump(queries_to_export, f, indent=4)
+        log.info("Saved %d queries to %s", len(queries_to_export), file)
+
+    else:
+    # import queries
+        with open(file, "r", encoding="UTF-8") as f:
+            try:
+                queries_to_import = json.load(f)
+            except json.decoder.JSONDecodeError:
+                log.error("Cannot parse file %s. A JSON file containing custom queries is required.", file)
+                sys.exit(1)
+
+        # compatibility import of legacy BloodHound's customqueries.json
+        if isinstance(queries_to_import, dict) and "queries" in queries_to_import:
+            log.warning("Detected legacy BloodHound's customqueries.json format. Some queries might not be compatible with BloodHound CE.")
+            compatible_queries = []
+            for entry in queries_to_import["queries"]:
+                if len(entry["queryList"]) > 1 or not entry["queryList"][0]["final"]:
+                    log.error('Query "%s" requires node selection, which is not compatible with BloodHound CE.', entry["name"])
+                    continue
+                compatible_queries.append({"name": entry["name"], "query": entry["queryList"][0]["query"]})
+            queries_to_import = compatible_queries
+
+        num_queries = 0
+        for query in queries_to_import:
+            try:
+                result = api.add_saved_query(query["name"], query["query"])
+                num_queries += 1
+            except ApiException as e:
+                if e.response.status_code == 400:
+                    log.error('Could not import query "%s": %s', query["name"], '\n'.join(error["message"] for error in e.response.json()["errors"]))
+                else:
+                    raise
+        log.info("Imported %d custom queries.", num_queries)