1"""
2Wrap the saltcheck module to copy files to ssh minion before running tests
3"""
4
5import logging
6import os
7import shutil
8import tarfile
9import tempfile
10from contextlib import closing
11
12import salt.utils.files
13import salt.utils.json
14import salt.utils.url
15
16log = logging.getLogger(__name__)
17
18
19def update_master_cache(states, saltenv="base"):
20    """
21    Replace standard saltcheck version with similar logic but replacing cp.cache_dir with
22        generating files, tar'ing them up, copying tarball to remote host, extracting tar
23        to state cache directory, and cleanup of files
24    """
25    cache = __opts__["cachedir"]
26    state_cache = os.path.join(cache, "files", saltenv)
27
28    # Setup for copying states to gendir
29    gendir = tempfile.mkdtemp()
30    trans_tar = salt.utils.files.mkstemp()
31    if "cp.fileclient_{}".format(id(__opts__)) not in __context__:
32        __context__[
33            "cp.fileclient_{}".format(id(__opts__))
34        ] = salt.fileclient.get_file_client(__opts__)
35
36    # generate cp.list_states output and save to gendir
37    cp_output = salt.utils.json.dumps(__salt__["cp.list_states"]())
38    cp_output_file = os.path.join(gendir, "cp_output.txt")
39    with salt.utils.files.fopen(cp_output_file, "w") as fp:
40        fp.write(cp_output)
41
42    # cp state directories to gendir
43    already_processed = []
44    sls_list = salt.utils.args.split_input(states)
45    for state_name in sls_list:
46        # generate low data for each state and save to gendir
47        state_low_file = os.path.join(gendir, state_name + ".low")
48        state_low_output = salt.utils.json.dumps(
49            __salt__["state.show_low_sls"](state_name)
50        )
51        with salt.utils.files.fopen(state_low_file, "w") as fp:
52            fp.write(state_low_output)
53
54        state_name = state_name.replace(".", os.sep)
55        if state_name in already_processed:
56            log.debug("Already cached state for %s", state_name)
57        else:
58            file_copy_file = os.path.join(gendir, state_name + ".copy")
59            log.debug("copying %s to %s", state_name, gendir)
60            qualified_name = salt.utils.url.create(state_name, saltenv)
61            # Duplicate cp.get_dir to gendir
62            copy_result = __context__["cp.fileclient_{}".format(id(__opts__))].get_dir(
63                qualified_name, gendir, saltenv
64            )
65            if copy_result:
66                copy_result = [dir.replace(gendir, state_cache) for dir in copy_result]
67                copy_result_output = salt.utils.json.dumps(copy_result)
68                with salt.utils.files.fopen(file_copy_file, "w") as fp:
69                    fp.write(copy_result_output)
70                already_processed.append(state_name)
71            else:
72                # If files were not copied, assume state.file.sls was given and just copy state
73                state_name = os.path.dirname(state_name)
74                file_copy_file = os.path.join(gendir, state_name + ".copy")
75                if state_name in already_processed:
76                    log.debug("Already cached state for %s", state_name)
77                else:
78                    qualified_name = salt.utils.url.create(state_name, saltenv)
79                    copy_result = __context__[
80                        "cp.fileclient_{}".format(id(__opts__))
81                    ].get_dir(qualified_name, gendir, saltenv)
82                    if copy_result:
83                        copy_result = [
84                            dir.replace(gendir, state_cache) for dir in copy_result
85                        ]
86                        copy_result_output = salt.utils.json.dumps(copy_result)
87                        with salt.utils.files.fopen(file_copy_file, "w") as fp:
88                            fp.write(copy_result_output)
89                        already_processed.append(state_name)
90
91    # turn gendir into tarball and remove gendir
92    try:
93        # cwd may not exist if it was removed but salt was run from it
94        cwd = os.getcwd()
95    except OSError:
96        cwd = None
97    os.chdir(gendir)
98    with closing(tarfile.open(trans_tar, "w:gz")) as tfp:
99        for root, dirs, files in salt.utils.path.os_walk(gendir):
100            for name in files:
101                full = os.path.join(root, name)
102                tfp.add(full[len(gendir) :].lstrip(os.sep))
103    if cwd:
104        os.chdir(cwd)
105    shutil.rmtree(gendir)
106
107    # Copy tarfile to ssh host
108    single = salt.client.ssh.Single(__opts__, "", **__salt__.kwargs)
109    thin_dir = __opts__["thin_dir"]
110    ret = single.shell.send(trans_tar, thin_dir)
111
112    # Clean up local tar
113    try:
114        os.remove(trans_tar)
115    except OSError:
116        pass
117
118    tar_path = os.path.join(thin_dir, os.path.basename(trans_tar))
119    # Extract remote tarball to cache directory and remove tar file
120    # TODO this could be better handled by a single state/connection due to ssh overhead
121    ret = __salt__["file.mkdir"](state_cache)
122    ret = __salt__["archive.tar"]("xf", tar_path, dest=state_cache)
123    ret = __salt__["file.remove"](tar_path)
124
125    return ret
126
127
128def run_state_tests(states, saltenv="base", check_all=False):
129    """
130    Define common functions to activite this wrapping module and tar copy.
131    After file copies are finished, run the usual local saltcheck function
132    """
133    ret = update_master_cache(states, saltenv)
134    ret = __salt__["saltcheck.run_state_tests_ssh"](
135        states, saltenv=saltenv, check_all=check_all
136    )
137    return ret
138
139
140def run_highstate_tests(saltenv="base"):
141    """
142    Lookup top files for minion, pass results to wrapped run_state_tests for copy and run
143    """
144    top_states = __salt__["state.show_top"]().get(saltenv)
145    state_string = ",".join(top_states)
146    ret = run_state_tests(state_string, saltenv)
147    return ret
148