481 lines
17 KiB
Python
481 lines
17 KiB
Python
# Copyright (C) 2022 Mohamed El Mouctar HAIDARA
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
from typing import Dict, Optional, Tuple, List, Set
|
|
|
|
from ansible.utils.display import Display
|
|
from graphviz import Digraph
|
|
|
|
from ansibleplaybookgrapher import PlaybookParser
|
|
from ansibleplaybookgrapher.graph import (
|
|
PlaybookNode,
|
|
RoleNode,
|
|
BlockNode,
|
|
Node,
|
|
PlayNode,
|
|
)
|
|
from ansibleplaybookgrapher.utils import get_play_colors, merge_dicts
|
|
|
|
display = Display()
|
|
|
|
# The supported protocol handlers to open roles and tasks from the viewer
|
|
OPEN_PROTOCOL_HANDLERS = {
|
|
"default": {"folder": "{path}", "file": "{path}"},
|
|
# https://code.visualstudio.com/docs/editor/command-line#_opening-vs-code-with-urls
|
|
"vscode": {
|
|
"folder": "vscode://file/{path}",
|
|
"file": "vscode://file/{path}:{line}:{column}",
|
|
},
|
|
# For custom, the formats need to be provided
|
|
"custom": {},
|
|
}
|
|
|
|
|
|
class Grapher:
|
|
def __init__(self, playbook_filenames: List[str]):
|
|
"""
|
|
|
|
:param playbook_filenames: List of playbooks to graph
|
|
"""
|
|
self.playbook_filenames = playbook_filenames
|
|
# Colors assigned to plays
|
|
|
|
self.plays_color = {}
|
|
# The usage of the roles in all playbooks
|
|
self.roles_usage: Dict["RoleNode", List[str]] = {}
|
|
|
|
# The parsed playbooks
|
|
self.playbook_nodes: List[PlaybookNode] = []
|
|
|
|
def parse(
|
|
self,
|
|
include_role_tasks: bool = False,
|
|
tags: List[str] = None,
|
|
skip_tags: List[str] = None,
|
|
group_roles_by_name: bool = False,
|
|
):
|
|
"""
|
|
Parses all the provided playbooks
|
|
:param include_role_tasks: Should we include the role tasks
|
|
:param tags: Only add plays and tasks tagged with these values
|
|
:param skip_tags: Only add plays and tasks whose tags do not match these values
|
|
:param group_roles_by_name: Group roles by name instead of considering them as separate nodes with different IDs
|
|
:return:
|
|
"""
|
|
|
|
for playbook_file in self.playbook_filenames:
|
|
display.display(f"Parsing playbook {playbook_file}")
|
|
parser = PlaybookParser(
|
|
tags=tags,
|
|
skip_tags=skip_tags,
|
|
playbook_filename=playbook_file,
|
|
include_role_tasks=include_role_tasks,
|
|
group_roles_by_name=group_roles_by_name,
|
|
)
|
|
playbook_node = parser.parse()
|
|
self.playbook_nodes.append(playbook_node)
|
|
|
|
# Setting colors for play
|
|
for play in playbook_node.plays:
|
|
# TODO: find a way to create visual distance between the generated colors
|
|
# https://stackoverflow.com/questions/9018016/how-to-compare-two-colors-for-similarity-difference
|
|
self.plays_color[play] = get_play_colors(play.id)
|
|
|
|
# Update the usage of the roles
|
|
self.roles_usage = merge_dicts(
|
|
self.roles_usage, playbook_node.roles_usage()
|
|
)
|
|
|
|
def graph(
|
|
self,
|
|
open_protocol_handler: str,
|
|
open_protocol_custom_formats: Dict[str, str] = None,
|
|
) -> Digraph:
|
|
"""
|
|
Generate the digraph graph
|
|
:param open_protocol_handler
|
|
:param open_protocol_custom_formats
|
|
:return:
|
|
"""
|
|
digraph = Digraph(
|
|
format="svg",
|
|
graph_attr=GraphvizGraphBuilder.DEFAULT_GRAPH_ATTR,
|
|
edge_attr=GraphvizGraphBuilder.DEFAULT_EDGE_ATTR,
|
|
)
|
|
# Map of the rules that have been built so far for all playbooks
|
|
roles_built = {}
|
|
for p in self.playbook_nodes:
|
|
builder = GraphvizGraphBuilder(
|
|
p,
|
|
digraph=digraph,
|
|
roles_usage=self.roles_usage,
|
|
roles_built=roles_built,
|
|
play_colors=self.plays_color,
|
|
open_protocol_handler=open_protocol_handler,
|
|
open_protocol_custom_formats=open_protocol_custom_formats,
|
|
)
|
|
builder.build_graphviz_graph()
|
|
roles_built.update(builder.roles_built)
|
|
|
|
return digraph
|
|
|
|
|
|
class GraphvizGraphBuilder:
|
|
"""
|
|
Build the graphviz graph
|
|
"""
|
|
|
|
DEFAULT_EDGE_ATTR = {"sep": "10", "esep": "5"}
|
|
DEFAULT_GRAPH_ATTR = {
|
|
"ratio": "fill",
|
|
"rankdir": "LR",
|
|
"concentrate": "true",
|
|
"ordering": "in",
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
playbook_node: PlaybookNode,
|
|
open_protocol_handler: str,
|
|
digraph: Digraph,
|
|
play_colors: Dict[PlayNode, Tuple[str, str]],
|
|
roles_usage: Dict[RoleNode, List[Node]] = None,
|
|
roles_built: Dict = None,
|
|
open_protocol_custom_formats: Dict[str, str] = None,
|
|
):
|
|
"""
|
|
|
|
:param playbook_node: Playbook parsed node
|
|
:param open_protocol_handler: The protocol handler name to use
|
|
:param digraph: Graphviz graph into which build the graph
|
|
:param open_protocol_custom_formats: The custom formats to use when the protocol handler is set to custom
|
|
"""
|
|
self.playbook_node = playbook_node
|
|
self.roles_usage = roles_usage or playbook_node.roles_usage()
|
|
self.play_colors = play_colors
|
|
# A map containing the roles that have been built so far
|
|
self.roles_built = roles_built or {}
|
|
|
|
self.open_protocol_handler = open_protocol_handler
|
|
# Merge the two dicts
|
|
formats = {**OPEN_PROTOCOL_HANDLERS, **{"custom": open_protocol_custom_formats}}
|
|
self.open_protocol_formats = formats[self.open_protocol_handler]
|
|
|
|
self.digraph = digraph
|
|
|
|
def build_node(
|
|
self,
|
|
graph: Digraph,
|
|
counter: int,
|
|
source: Node,
|
|
destination: Node,
|
|
color: str,
|
|
shape: str = "octagon",
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Render a generic node in the graph
|
|
:param graph: The graph to render the node to
|
|
:param source: The source node
|
|
:param destination: The RoleNode to render
|
|
:param color: The color to apply
|
|
:param counter: The counter for this node
|
|
:param shape: the default shape of the node
|
|
:return:
|
|
"""
|
|
|
|
node_label_prefix = kwargs.get("node_label_prefix", "")
|
|
|
|
if isinstance(destination, BlockNode):
|
|
self.build_block(
|
|
graph,
|
|
counter,
|
|
source=source,
|
|
destination=destination,
|
|
color=color,
|
|
)
|
|
elif isinstance(destination, RoleNode):
|
|
self.build_role(
|
|
graph,
|
|
counter,
|
|
source=source,
|
|
destination=destination,
|
|
color=color,
|
|
)
|
|
else:
|
|
# Here we have a TaskNode
|
|
edge_label = f"{counter} {destination.when}"
|
|
# Task node
|
|
graph.node(
|
|
destination.id,
|
|
label=node_label_prefix + destination.name,
|
|
shape=shape,
|
|
id=destination.id,
|
|
tooltip=destination.name,
|
|
color=color,
|
|
URL=self.get_node_url(destination, "file"),
|
|
)
|
|
# Edge from parent to task
|
|
graph.edge(
|
|
source.id,
|
|
destination.id,
|
|
label=edge_label,
|
|
color=color,
|
|
fontcolor=color,
|
|
id=f"edge_{counter}_{source.id}_{destination.id}",
|
|
tooltip=edge_label,
|
|
labeltooltip=edge_label,
|
|
)
|
|
|
|
def build_block(
|
|
self,
|
|
graph: Digraph,
|
|
counter: int,
|
|
source: Node,
|
|
destination: BlockNode,
|
|
color: str,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Render a block in the graph.
|
|
A BlockNode is a special node: a cluster is created instead of a normal node.
|
|
:param graph: The graph to render the block into
|
|
:param counter: The counter for this block in the graph
|
|
:param source: The source node
|
|
:param destination: The BlockNode to render
|
|
:param color: The color to apply
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
edge_label = f"{counter}"
|
|
|
|
# Edge from parent to the block node inside the cluster
|
|
graph.edge(
|
|
source.id,
|
|
destination.id,
|
|
label=edge_label,
|
|
color=color,
|
|
fontcolor=color,
|
|
tooltip=edge_label,
|
|
id=f"edge_{counter}_{source.id}_{destination.id}",
|
|
labeltooltip=edge_label,
|
|
)
|
|
|
|
# BlockNode is a special node: a cluster is created instead of a normal node
|
|
with graph.subgraph(name=f"cluster_{destination.id}") as cluster_block_subgraph:
|
|
# block node
|
|
cluster_block_subgraph.node(
|
|
destination.id,
|
|
label=f"[block] {destination.name}",
|
|
shape="box",
|
|
id=destination.id,
|
|
tooltip=destination.name,
|
|
color=color,
|
|
labeltooltip=destination.name,
|
|
URL=self.get_node_url(destination, "file"),
|
|
)
|
|
|
|
# The reverse here is a little hack due to how graphviz render nodes inside a cluster by reversing them.
|
|
# Don't really know why for the moment neither if there is an attribute to change that.
|
|
for b_counter, task in enumerate(reversed(destination.tasks)):
|
|
self.build_node(
|
|
cluster_block_subgraph,
|
|
source=destination,
|
|
destination=task,
|
|
counter=len(destination.tasks) - b_counter,
|
|
color=color,
|
|
)
|
|
|
|
def build_role(
|
|
self,
|
|
graph: Digraph,
|
|
counter: int,
|
|
source: Node,
|
|
destination: RoleNode,
|
|
color: str,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Render a role in the graph
|
|
:param graph: The graph to render the role into
|
|
:param counter: The counter for this role in the graph
|
|
:param source: The source node
|
|
:param destination: The RoleNode to render
|
|
:param color: The color to apply
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
|
|
if destination.include_role: # For include_role, we point to a file
|
|
url = self.get_node_url(destination, "file")
|
|
else: # For normal role invocation, we point to the folder
|
|
url = self.get_node_url(destination, "folder")
|
|
|
|
role_edge_label = f"{counter} {destination.when}"
|
|
|
|
# from parent to the role node
|
|
graph.edge(
|
|
source.id,
|
|
destination.id,
|
|
label=role_edge_label,
|
|
color=color,
|
|
fontcolor=color,
|
|
id=f"edge_{counter}_{source.id}_{destination.id}",
|
|
tooltip=role_edge_label,
|
|
labeltooltip=role_edge_label,
|
|
)
|
|
|
|
# check if we already built this role
|
|
role_to_render = self.roles_built.get(destination.id, None)
|
|
if role_to_render is None:
|
|
# Merge the colors for each play where this role is used
|
|
role_plays = self.roles_usage[destination]
|
|
# Graphviz support providing multiple colors separated by :
|
|
if len(role_plays) > 1:
|
|
# If the role is used in multiple plays, we take black as the default color
|
|
role_color = "black"
|
|
else:
|
|
colors = list(map(self.play_colors.get, role_plays))[0]
|
|
role_color = colors[0]
|
|
|
|
self.roles_built[destination.id] = destination
|
|
|
|
with graph.subgraph(name=destination.name, node_attr={}) as role_subgraph:
|
|
role_subgraph.node(
|
|
destination.id,
|
|
id=destination.id,
|
|
label=f"[role] {destination.name}",
|
|
tooltip=destination.name,
|
|
color=color,
|
|
URL=url,
|
|
)
|
|
# role tasks
|
|
for role_task_counter, role_task in enumerate(destination.tasks, 1):
|
|
self.build_node(
|
|
role_subgraph,
|
|
source=destination,
|
|
destination=role_task,
|
|
counter=role_task_counter,
|
|
color=role_color,
|
|
)
|
|
|
|
def build_graphviz_graph(self):
|
|
"""
|
|
Convert the PlaybookNode to the graphviz dot format
|
|
:return:
|
|
"""
|
|
display.vvv(f"Converting the graph to the dot format for graphviz")
|
|
# root node
|
|
self.digraph.node(
|
|
self.playbook_node.name,
|
|
style="dotted",
|
|
id=self.playbook_node.id,
|
|
URL=self.get_node_url(self.playbook_node, "file"),
|
|
)
|
|
|
|
for play_counter, play in enumerate(self.playbook_node.plays, 1):
|
|
with self.digraph.subgraph(name=play.name) as play_subgraph:
|
|
color, play_font_color = self.play_colors[play]
|
|
play_tooltip = (
|
|
",".join(play.hosts) if len(play.hosts) > 0 else play.name
|
|
)
|
|
|
|
# play node
|
|
play_subgraph.node(
|
|
play.id,
|
|
id=play.id,
|
|
label=play.name,
|
|
style="filled",
|
|
shape="box",
|
|
color=color,
|
|
fontcolor=play_font_color,
|
|
tooltip=play_tooltip,
|
|
URL=self.get_node_url(play, "file"),
|
|
)
|
|
|
|
# edge from root node to play
|
|
playbook_to_play_label = f"{play_counter} {play.name}"
|
|
self.digraph.edge(
|
|
self.playbook_node.name,
|
|
play.id,
|
|
id=f"edge_{self.playbook_node.id}_{play.id}",
|
|
label=playbook_to_play_label,
|
|
color=color,
|
|
fontcolor=color,
|
|
tooltip=playbook_to_play_label,
|
|
labeltooltip=playbook_to_play_label,
|
|
)
|
|
|
|
# pre_tasks
|
|
for pre_task_counter, pre_task in enumerate(play.pre_tasks, 1):
|
|
self.build_node(
|
|
play_subgraph,
|
|
counter=pre_task_counter,
|
|
source=play,
|
|
destination=pre_task,
|
|
color=color,
|
|
node_label_prefix="[pre_task] ",
|
|
)
|
|
|
|
# roles
|
|
for role_counter, role in enumerate(play.roles, 1):
|
|
self.build_role(
|
|
play_subgraph,
|
|
source=play,
|
|
destination=role,
|
|
counter=role_counter + len(play.pre_tasks),
|
|
color=color,
|
|
)
|
|
|
|
# tasks
|
|
for task_counter, task in enumerate(play.tasks, 1):
|
|
self.build_node(
|
|
play_subgraph,
|
|
source=play,
|
|
destination=task,
|
|
counter=len(play.pre_tasks) + len(play.roles) + task_counter,
|
|
color=color,
|
|
node_label_prefix="[task] ",
|
|
)
|
|
|
|
# post_tasks
|
|
for post_task_counter, post_task in enumerate(play.post_tasks, 1):
|
|
self.build_node(
|
|
play_subgraph,
|
|
source=play,
|
|
destination=post_task,
|
|
counter=len(play.pre_tasks)
|
|
+ len(play.roles)
|
|
+ len(play.tasks)
|
|
+ post_task_counter,
|
|
color=color,
|
|
node_label_prefix="[post_task] ",
|
|
)
|
|
|
|
def get_node_url(self, node: Node, node_type: str) -> Optional[str]:
|
|
"""
|
|
Get the node url based on the chosen protocol
|
|
:param node_type: task or role
|
|
:param node: the node to get the url for
|
|
:return:
|
|
"""
|
|
if node.path:
|
|
url = self.open_protocol_formats[node_type].format(
|
|
path=node.path, line=node.line, column=node.column
|
|
)
|
|
display.vvvv(f"Open protocol URL for node {node}: {url}")
|
|
return url
|
|
|
|
return None
|