1""" 2SaltStack Extend 3~~~~~~~~~~~~~~~~ 4 5A templating tool for extending SaltStack. 6 7Takes a template directory and merges it into a SaltStack source code 8directory. This tool uses Jinja2 for templating. 9 10This tool is accessed using `salt-extend` 11 12 :codeauthor: Anthony Shaw <anthonyshaw@apache.org> 13""" 14 15 16import logging 17import os 18import shutil 19import sys 20import tempfile 21from datetime import date 22 23import salt.utils.files 24import salt.version 25from jinja2 import Template 26from salt.serializers.yaml import deserialize 27from salt.utils.odict import OrderedDict 28 29log = logging.getLogger(__name__) 30 31try: 32 import click 33 34 HAS_CLICK = True 35except ImportError as ie: 36 HAS_CLICK = False 37 38TEMPLATE_FILE_NAME = "template.yml" 39 40 41def _get_template(path, option_key): 42 """ 43 Get the contents of a template file and provide it as a module type 44 45 :param path: path to the template.yml file 46 :type path: ``str`` 47 48 :param option_key: The unique key of this template 49 :type option_key: ``str`` 50 51 :returns: Details about the template 52 :rtype: ``tuple`` 53 """ 54 with salt.utils.files.fopen(path, "r") as template_f: 55 template = deserialize(template_f) 56 info = (option_key, template.get("description", ""), template) 57 return info 58 59 60def _fetch_templates(src): 61 """ 62 Fetch all of the templates in the src directory 63 64 :param src: The source path 65 :type src: ``str`` 66 67 :rtype: ``list`` of ``tuple`` 68 :returns: ``list`` of ('key', 'description') 69 """ 70 templates = [] 71 log.debug("Listing contents of %s", src) 72 for item in os.listdir(src): 73 s = os.path.join(src, item) 74 if os.path.isdir(s): 75 template_path = os.path.join(s, TEMPLATE_FILE_NAME) 76 if os.path.isfile(template_path): 77 templates.append(_get_template(template_path, item)) 78 else: 79 log.debug( 80 "Directory does not contain %s %s", 81 template_path, 82 TEMPLATE_FILE_NAME, 83 ) 84 return templates 85 86 87def _mergetree(src, dst): 88 """ 89 Akin to shutils.copytree but over existing directories, does a recursive merge copy. 90 91 :param src: The source path 92 :type src: ``str`` 93 94 :param dst: The destination path 95 :type dst: ``str`` 96 """ 97 for item in os.listdir(src): 98 s = os.path.join(src, item) 99 d = os.path.join(dst, item) 100 if os.path.isdir(s): 101 log.info("Copying folder %s to %s", s, d) 102 if os.path.exists(d): 103 _mergetree(s, d) 104 else: 105 shutil.copytree(s, d) 106 else: 107 log.info("Copying file %s to %s", s, d) 108 shutil.copy2(s, d) 109 110 111def _mergetreejinja(src, dst, context): 112 """ 113 Merge directory A to directory B, apply Jinja2 templating to both 114 the file/folder names AND to the contents of the files 115 116 :param src: The source path 117 :type src: ``str`` 118 119 :param dst: The destination path 120 :type dst: ``str`` 121 122 :param context: The dictionary to inject into the Jinja template as context 123 :type context: ``dict`` 124 """ 125 for item in os.listdir(src): 126 s = os.path.join(src, item) 127 d = os.path.join(dst, item) 128 if os.path.isdir(s): 129 log.info("Copying folder %s to %s", s, d) 130 if os.path.exists(d): 131 _mergetreejinja(s, d, context) 132 else: 133 os.mkdir(d) 134 _mergetreejinja(s, d, context) 135 else: 136 if item != TEMPLATE_FILE_NAME: 137 d = Template(d).render(context) 138 log.info("Copying file %s to %s", s, d) 139 with salt.utils.files.fopen(s, "r") as source_file: 140 src_contents = salt.utils.stringutils.to_unicode(source_file.read()) 141 dest_contents = Template(src_contents).render(context) 142 with salt.utils.files.fopen(d, "w") as dest_file: 143 dest_file.write(salt.utils.stringutils.to_str(dest_contents)) 144 145 146def _prompt_user_variable(var_name, default_value): 147 """ 148 Prompt the user to enter the value of a variable 149 150 :param var_name: The question to ask the user 151 :type var_name: ``str`` 152 153 :param default_value: The default value 154 :type default_value: ``str`` 155 156 :rtype: ``str`` 157 :returns: the value from the user 158 """ 159 return click.prompt(var_name, default=default_value) 160 161 162def _prompt_choice(var_name, options): 163 """ 164 Prompt the user to choose between a list of options, index each one by adding an enumerator 165 based on https://github.com/audreyr/cookiecutter/blob/master/cookiecutter/prompt.py#L51 166 167 :param var_name: The question to ask the user 168 :type var_name: ``str`` 169 170 :param options: A list of options 171 :type options: ``list`` of ``tupple`` 172 173 :rtype: ``tuple`` 174 :returns: The selected user 175 """ 176 choice_map = OrderedDict( 177 ("{}".format(i), value) 178 for i, value in enumerate(options, 1) 179 if value[0] != "test" 180 ) 181 choices = choice_map.keys() 182 default = "1" 183 184 choice_lines = [ 185 "{} - {} - {}".format(c[0], c[1][0], c[1][1]) for c in choice_map.items() 186 ] 187 prompt = "\n".join( 188 ( 189 "Select {}:".format(var_name), 190 "\n".join(choice_lines), 191 "Choose from {}".format(", ".join(choices)), 192 ) 193 ) 194 195 user_choice = click.prompt(prompt, type=click.Choice(choices), default=default) 196 return choice_map[user_choice] 197 198 199def apply_template(template_dir, output_dir, context): 200 """ 201 Apply the template from the template directory to the output 202 using the supplied context dict. 203 204 :param src: The source path 205 :type src: ``str`` 206 207 :param dst: The destination path 208 :type dst: ``str`` 209 210 :param context: The dictionary to inject into the Jinja template as context 211 :type context: ``dict`` 212 """ 213 _mergetreejinja(template_dir, output_dir, context) 214 215 216def run( 217 extension=None, 218 name=None, 219 description=None, 220 salt_dir=None, 221 merge=False, 222 temp_dir=None, 223): 224 """ 225 A template factory for extending the salt ecosystem 226 227 :param extension: The extension type, e.g. 'module', 'state', if omitted, user will be prompted 228 :type extension: ``str`` 229 230 :param name: Python-friendly name for the module, if omitted, user will be prompted 231 :type name: ``str`` 232 233 :param description: A description of the extension, if omitted, user will be prompted 234 :type description: ``str`` 235 236 :param salt_dir: The targeted Salt source directory 237 :type salt_dir: ``str`` 238 239 :param merge: Merge with salt directory, `False` to keep separate, `True` to merge trees. 240 :type merge: ``bool`` 241 242 :param temp_dir: The directory for generated code, if omitted, system temp will be used 243 :type temp_dir: ``str`` 244 """ 245 if not HAS_CLICK: 246 print("click is not installed, please install using pip") 247 sys.exit(1) 248 249 if salt_dir is None: 250 salt_dir = "." 251 252 MODULE_OPTIONS = _fetch_templates(os.path.join(salt_dir, "templates")) 253 254 if extension is None: 255 print("Choose which type of extension you are developing for SaltStack") 256 extension_type = "Extension type" 257 chosen_extension = _prompt_choice(extension_type, MODULE_OPTIONS) 258 else: 259 if extension not in list(zip(*MODULE_OPTIONS))[0]: 260 print("Module extension option not valid") 261 sys.exit(1) 262 263 chosen_extension = [m for m in MODULE_OPTIONS if m[0] == extension][0] 264 265 extension_type = chosen_extension[0] 266 extension_context = chosen_extension[2] 267 268 if name is None: 269 print("Enter the short name for the module (e.g. mymodule)") 270 name = _prompt_user_variable("Module name", "") 271 272 if description is None: 273 description = _prompt_user_variable("Short description of the module", "") 274 275 template_dir = "templates/{}".format(extension_type) 276 module_name = name 277 278 param_dict = { 279 "version": salt.version.SaltStackVersion.next_release().name, 280 "module_name": module_name, 281 "short_description": description, 282 "release_date": date.today().strftime("%Y-%m-%d"), 283 "year": date.today().strftime("%Y"), 284 } 285 286 # get additional questions from template 287 additional_context = {} 288 for key, val in extension_context.get("questions", {}).items(): 289 # allow templates to be used in default values. 290 default = Template(val.get("default", "")).render(param_dict) 291 292 prompt_var = _prompt_user_variable(val["question"], default) 293 additional_context[key] = prompt_var 294 295 context = param_dict.copy() 296 context.update(extension_context) 297 context.update(additional_context) 298 299 if temp_dir is None: 300 temp_dir = tempfile.mkdtemp() 301 302 apply_template(template_dir, temp_dir, context) 303 304 if not merge: 305 path = temp_dir 306 else: 307 _mergetree(temp_dir, salt_dir) 308 path = salt_dir 309 310 log.info("New module stored in %s", path) 311 return path 312 313 314if __name__ == "__main__": 315 run() 316