From e3584bdc1680cebb2bd236864cd0e29b608129f8 Mon Sep 17 00:00:00 2001 From: stg-annon <14135675+stg-annon@users.noreply.github.com> Date: Sat, 15 Jan 2022 18:09:33 -0500 Subject: [PATCH] Update tagGraph plugin to Script (#32) * init tagGraph * Update tagGraph plugin to Script removes tagGraph as a plugin and instead runs as a script allowing it to work on remote and Docker instances -need for use of sqlite to fetch graph data +minimal documentation +stash like theme to graph --- scripts/tagGraph/README.md | 18 ++++ scripts/tagGraph/requirements.txt | 2 + scripts/tagGraph/stash_interface.py | 144 ++++++++++++++++++++++++++++ scripts/tagGraph/tag_graph.py | 69 +++++++++++++ 4 files changed, 233 insertions(+) create mode 100644 scripts/tagGraph/README.md create mode 100644 scripts/tagGraph/requirements.txt create mode 100644 scripts/tagGraph/stash_interface.py create mode 100644 scripts/tagGraph/tag_graph.py diff --git a/scripts/tagGraph/README.md b/scripts/tagGraph/README.md new file mode 100644 index 0000000..b66f766 --- /dev/null +++ b/scripts/tagGraph/README.md @@ -0,0 +1,18 @@ + +# Tag Graph Generator + +## Requirements + * python >= 3.7.X + * `pip install -r requirements.txt` + +## Usage + +ensure `STASH_SETTINGS` is configured properly +> **⚠️ Note:** if you are connecting to a remote/docker instance of stash you will need to change this + +run `python .\tag_graph.py` + +## Customizing the graph +set `SHOW_OPTIONS` to `True` and you will get an interface to play around with that will affect what the graph looks like. + +for more info see [pyvis docs](https://pyvis.readthedocs.io/en/latest/tutorial.html#using-the-configuration-ui-to-dynamically-tweak-network-settings) diff --git a/scripts/tagGraph/requirements.txt b/scripts/tagGraph/requirements.txt new file mode 100644 index 0000000..2e17f0e --- /dev/null +++ b/scripts/tagGraph/requirements.txt @@ -0,0 +1,2 @@ +pyvis==0.1.9 +requests==2.25.1 \ No newline at end of file diff --git a/scripts/tagGraph/stash_interface.py b/scripts/tagGraph/stash_interface.py new file mode 100644 index 0000000..003bac6 --- /dev/null +++ b/scripts/tagGraph/stash_interface.py @@ -0,0 +1,144 @@ +import re, sys, requests + +class StashInterface: + port = "" + url = "" + headers = { + "Accept-Encoding": "gzip, deflate", + "Content-Type": "application/json", + "Accept": "application/json", + "Connection": "keep-alive", + "DNT": "1" + } + cookies = {} + + def __init__(self, conn, fragments={}): + global log + + if conn.get("Logger"): + log = conn.get("Logger") + else: + raise Exception("No logger passed to StashInterface") + + self.port = conn['Port'] if conn.get('Port') else '9999' + scheme = conn['Scheme'] if conn.get('Scheme') else 'http' + + # Session cookie for authentication + + self.cookies = {} + if conn.get('SessionCookie'): + self.cookies.update({ + 'session': conn['SessionCookie']['Value'] + }) + + domain = conn['Domain'] if conn.get('Domain') else 'localhost' + + # Stash GraphQL endpoint + self.url = f'{scheme}://{domain}:{self.port}/graphql' + + try: + self.get_stash_config() + except Exception: + log.error(f"Could not connect to Stash at {self.url}") + sys.exit() + + log.info(f"Using Stash's GraphQl endpoint at {self.url}") + + self.fragments = fragments + + def __resolveFragments(self, query): + + fragmentReferences = list(set(re.findall(r'(?<=\.\.\.)\w+', query))) + fragments = [] + for ref in fragmentReferences: + fragments.append({ + "fragment": ref, + "defined": bool(re.search("fragment {}".format(ref), query)) + }) + + if all([f["defined"] for f in fragments]): + return query + else: + for fragment in [f["fragment"] for f in fragments if not f["defined"]]: + if fragment not in self.fragments: + raise Exception(f'GraphQL error: fragment "{fragment}" not defined') + query += self.fragments[fragment] + return self.__resolveFragments(query) + + def __callGraphQL(self, query, variables=None): + + query = self.__resolveFragments(query) + + json_request = {'query': query} + if variables is not None: + json_request['variables'] = variables + + response = requests.post(self.url, json=json_request, headers=self.headers, cookies=self.cookies) + + if response.status_code == 200: + result = response.json() + + if result.get("errors"): + for error in result["errors"]: + log.error(f"GraphQL error: {error}") + if result.get("error"): + for error in result["error"]["errors"]: + log.error(f"GraphQL error: {error}") + if result.get("data"): + return result['data'] + elif response.status_code == 401: + sys.exit("HTTP Error 401, Unauthorized. Cookie authentication most likely failed") + else: + raise ConnectionError( + "GraphQL query failed:{} - {}. Query: {}. Variables: {}".format( + response.status_code, response.content, query, variables) + ) + + def __match_alias_item(self, search, items): + item_matches = {} + for item in items: + if re.match(rf'{search}$', item.name, re.IGNORECASE): + log.debug(f'matched "{search}" to "{item.name}" ({item.id}) using primary name') + item_matches[item.id] = item + if not item.aliases: + continue + for alias in item.aliases: + if re.match(rf'{search}$', alias.strip(), re.IGNORECASE): + log.debug(f'matched "{search}" to "{alias}" ({item.id}) using alias') + item_matches[item.id] = item + return list(item_matches.values()) + + def get_stash_config(self): + query = """ + query Configuration { + configuration { general { stashes{ path } } } + } + """ + result = self.__callGraphQL(query) + return result['configuration'] + + def get_tags_with_relations(self): + query = """ + query FindTags($filter: FindFilterType, $tag_filter: TagFilterType) { + findTags(filter: $filter, tag_filter: $tag_filter) { + count + tags { + id + name + parents { id } + children { id } + } + } + } + """ + + variables = { + "tag_filter":{ + "child_count":{"modifier": "GREATER_THAN", "value": 0}, + "OR": { + "parent_count": {"modifier": "GREATER_THAN", "value": 0}} + }, + "filter": {"q":"", "per_page":-1} + } + result = self.__callGraphQL(query, variables) + return result['findTags']['tags'] diff --git a/scripts/tagGraph/tag_graph.py b/scripts/tagGraph/tag_graph.py new file mode 100644 index 0000000..bc66dac --- /dev/null +++ b/scripts/tagGraph/tag_graph.py @@ -0,0 +1,69 @@ + +import os, sys, json +import logging as log + +# local dependencies +from stash_interface import StashInterface +# external dependencies +from pyvis.network import Network + + +### USER CONFIG ### +STASH_SETTINGS = { + "Scheme":"http", + "Domain": "localhost", + "Port": "9999", + "Logger": log, +} +SHOW_OPTIONS = False + + +def main(): + global stash + + log.basicConfig(level=log.INFO, format='%(levelname)s: %(message)s') + + stash = StashInterface(STASH_SETTINGS) + + log.info("getting tags from stash...") + tags = stash.get_tags_with_relations() + + log.info("generating graph...") + + + if SHOW_OPTIONS: + G = Network(directed=True, height="100%", width="66%", bgcolor="#202b33", font_color="white") + G.show_buttons() + else: + G = Network(directed=True, height="100%", width="100%", bgcolor="#202b33", font_color="white") + + node_theme = { + "border": "#adb5bd", + "background":"#394b59", + "highlight":{ + "border": "#137cbd", + "background":"#FFFFFF" + } + } + edge_theme = { + "color": "#FFFFFF", + "highlight":"#137cbd" + } + + # create all nodes + for tag in tags: + G.add_node(tag["id"], label=tag["name"], color=node_theme ) + # create all edges + for tag in tags: + for child in tag["children"]: + G.add_edge( tag["id"], child["id"], color=edge_theme ) + + + current_abs_path = os.path.dirname(os.path.abspath(__file__)) + save_path = os.path.join(current_abs_path, "tag_graph.html") + + G.save_graph(save_path) + log.info(f'saved graph to "{save_path}"') + +if __name__ == '__main__': + main() \ No newline at end of file