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