1#!/usr/bin/env python3
2from pathlib import Path
3from pprint import pformat
4from re import compile as regex
5from typing import Optional, Dict, List, Set, Any, Union
6
7import click
8from rich import print
9from rich.text import Text
10from rich.tree import Tree
11from yaml import safe_load
12
13
14def print_help() -> None:
15 ctx = click.get_current_context()
16 click.echo(ctx.get_help())
17 ctx.exit()
18
19
20class Role:
21 def __init__(self, role_dir: Path, root_dir: Path) -> None:
22 self.role_name = str(role_dir.relative_to(root_dir))
23 self.role_dir = role_dir.resolve()
24 self.dependencies: Dict[str, List[str]] = self.scan_dependencies()
25 # TODO: find tags
26
27 @property
28 def dependency_file(self) -> Path:
29 return self.role_dir / "meta" / "main.yaml"
30
31 def scan_dependencies(self) -> List[str]:
32 result = {}
33 if self.dependency_file.is_file():
34 meta = safe_load(self.dependency_file.open())
35 if (
36 meta is not None
37 and "dependencies" in meta
38 and meta["dependencies"] is not None
39 ):
40 for role in meta["dependencies"]:
41 if type(role) == str:
42 result[role] = []
43 else:
44 result[role["role"]] = role["tags"] if "tags" in role else []
45 return result
46
47 def tree(
48 self, roles: Dict[str, "Role"], print_tags: bool = False, tags: List[str] = []
49 ):
50 root = Tree(
51 f'{self.role_name}{" " + str(tags) if tags and print_tags else ""}',
52 style="green" if not tags else "bold blue",
53 )
54 for dependency, tags in self.dependencies.items():
55 if dependency not in roles:
56 root.add(
57 Text(
58 f'{dependency}{tags if tags and print_tags else ""}',
59 style="bold red",
60 )
61 )
62 else:
63 root.add(
64 roles[dependency].tree(roles, print_tags=print_tags, tags=tags)
65 )
66 return root
67
68
69def recurse(
70 element: Union[Dict[str, Any], List[Any]],
71 collection: Optional[Dict[str, List[str]]] = None,
72) -> None:
73 collection = collection if collection is not None else {}
74 if type(element) is dict:
75 for key, value in element.items():
76 if key == "include_role":
77 collection[value["name"]] = element["tags"] if "tags" in element else []
78 if key == "roles":
79 for role in value:
80 if type(role) == str:
81 collection[role] = []
82 else:
83 collection[role["role"]] = (
84 role["tags"] if "tags" in role else []
85 )
86 else:
87 recurse(value, collection)
88 elif type(element) is list:
89 for value in element:
90 recurse(value, collection)
91 return collection
92
93
94def get_roles_from_playbook(playbook: Path) -> Dict[str, List[str]]:
95 play = safe_load(playbook.open())
96 return recurse(play)
97
98
99class DependencyTree:
100 def __init__(self, root_dir: Path, playbook: Optional[Path] = None) -> None:
101 self.root_dir: Path = root_dir
102 self.playbook = playbook
103 self.playbook_roles = (
104 get_roles_from_playbook(playbook) if self.playbook else None
105 )
106 self.roles: Dict[str, Role] = {}
107 self.scan()
108
109 @staticmethod
110 def _is_role(candidate: Path) -> bool:
111 if not candidate.is_dir():
112 return False
113 yaml_files = set(str(c) for c in candidate.rglob("main.yaml"))
114 return (
115 str(candidate / "tasks" / "main.yaml") in yaml_files
116 or str(candidate / "meta" / "main.yaml") in yaml_files
117 )
118
119 def scan(self) -> None:
120 for path in self.root_dir.glob("**/*"):
121 if self._is_role(path):
122 self.roles[str(path.relative_to(self.root_dir))] = Role(
123 path, self.root_dir
124 )
125
126 def print(
127 self,
128 root_text: str = ".",
129 filter_roles: Optional[Set[str]] = None,
130 tags: bool = False,
131 ) -> None:
132 root = Tree(Text(root_text))
133 for name, role in self.roles.items():
134 if (not filter_roles or role.role_name in filter_roles) and (
135 not self.playbook or role.role_name in self.playbook_roles
136 ):
137 root.add(
138 role.tree(
139 self.roles,
140 print_tags=tags,
141 tags=self.playbook_roles[name]
142 if self.playbook_roles and name in self.playbook_roles
143 else [],
144 )
145 )
146 print(root)
147
148
149def generate_roles_tree(
150 root_dir: Path, playbook: Optional[Path] = None
151) -> DependencyTree:
152 return DependencyTree(root_dir, playbook)
153
154
155ROLES_PATH_PAT = regex(r"^roles_path\s*\=\s*(.+)$")
156
157
158def get_roles_path_from_ansible_cfg(ansible_cfg: Path) -> Path:
159 for line in ansible_cfg.open().readlines():
160 if match := ROLES_PATH_PAT.match(line):
161 return Path(match.group(1))
162
163
164@click.group()
165def cli():
166 pass
167
168
169@cli.command()
170@click.option(
171 "-r",
172 "--role",
173 default=None,
174 help="Name of role for which to generate a dependency tree for",
175)
176@click.option(
177 "-d",
178 "--role-dir",
179 default=None,
180 help="Directory to search for roles in (generates dependency tree for all roles if no playbook [-p] or specific role [-r] is provided)",
181)
182@click.option(
183 "-p",
184 "--playbook",
185 default=None,
186 help="Path to a playbook YAML file from which a dependency tree will be generated",
187)
188@click.option(
189 "-f",
190 "--find",
191 default=None,
192 help="Find a particular role in the dependency tree and highlight the path to it (not implemented)",
193)
194@click.option("--tags", is_flag=True, help="Display tags in the tree as well")
195def roles(
196 role: Optional[str],
197 role_dir: Optional[str],
198 playbook: Optional[str],
199 find: Optional[str],
200 tags: bool = False,
201) -> None:
202 # TODO: find
203 if role_dir is None:
204 if (Path.cwd() / "ansible.cfg").is_file():
205 role_dir = get_roles_path_from_ansible_cfg(Path.cwd() / "ansible.cfg")
206 elif (Path.cwd() / "roles").is_dir():
207 role_dir = Path.cwd() / "roles"
208 else:
209 role_dir = Path.cwd()
210 elif Path(role_dir).is_dir():
211 role_dir = Path(role_dir)
212 else:
213 print(f"Role directory: {role_dir} is not accessible")
214 sys.exit(1)
215 if playbook is not None:
216 if not Path(playbook).is_file():
217 print(f"Playbook file: {playbook} is not accessible.")
218 sys.exit(1)
219 tree = generate_roles_tree(role_dir, Path(playbook))
220 else:
221 tree = generate_roles_tree(role_dir)
222 tree.print(
223 root_text=str(role_dir), filter_roles={role} if role else None, tags=tags
224 )
225
226
227if __name__ == "__main__":
228 cli()