<commit_before>import salt.client
import salt.config
import lockfile
import os
import os.path
import random
import shutil
import string
import simplejson as json
import yaml
try:
    from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper

import util

__author__ = 'bkeroack'

class FatalSaltError(Exception):
    pass

class OSTypes:
    Windows = 0
    Unix_like = 1

class RemoteCommands:
    def __init__(self, sc):
        self.sc = sc

    def get_os(self, server):
        resp = self.sc.salt_command(server, 'grains.item', ["os"])
        return OSTypes.Windows if resp[server]['os'] == "Windows" else OSTypes.Unix_like

    def create_directory(self, server_list, path):
        assert isinstance(server_list, list)
        return self.sc.salt_command(server_list, 'file.mkdir', [path], opts={'expr_form': 'list'})

    def delete_directory_win(self, server_list, path):
        return self.sc.salt_command(server_list, 'cmd.run', ["Remove-Item -Recurse -Force {}".format(path)],
                                    opts={'expr_form': 'list',
                                          'shell': 'powershell'})

    def delete_directory_unix(self, server_list, path):
        return self.sc.salt_command(server_list, 'cmd.run', ["rm -rf {}".format(path)], opts={'expr_form': 'list'})

    def delete_directory(self, server_list, path):
        assert isinstance(server_list, list)
        res = list()
        win_s = list()
        unix_s = list()
        for s in server_list:
            ost = self.get_os(s)
            if ost == OSTypes.Windows:
                win_s.append(s)
            elif os == OSTypes.Unix_like:
                unix_s.append(s)
        if len(win_s) > 0:
            res.append(self.delete_directory_win(win_s, path))
        if len(unix_s) > 0:
            res.append(self.delete_directory_unix(unix_s, path))
        return res

    def push_file(self, target, local_path, remote_path):
        r_postfix = ''.join(random.sample(string.letters + string.digits, 20))
        root = self.sc.sls_dir
        newname = os.path.basename(local_path) + r_postfix
        new_fullpath = "{}/{}".format(root, newname)
        shutil.copy(local_path, new_fullpath)
        salt_uri = 'salt://{}/{}'.format(self.sc.settings['daft.salt.slsdir'], newname)
        util.debugLog(self, "push_file: salt_uri: {}".format(salt_uri))
        util.debugLog(self, "push_file: remote_path: {}".format(remote_path))
        res = self.sc.salt_command(target, 'cp.get_file', [salt_uri, remote_path])
        util.debugLog(self, "push_file: resp: {}".format(res))
        os.unlink(new_fullpath)
        return {'success': {'target': target, 'remote_path': remote_path}}

    def push_key(self, server_list, file, name, ext):
        res = list()
        for s in server_list:
            ost = self.get_os(s)
            results = {
                'server': s,
                'os': ost,
                'results': []
            }
            if ost == OSTypes.Windows:
                path = 'C:\Program Files (x86)\Git\.ssh'
                fullpath = path + "\{}{}".format("id_rsa", ext)
            elif ost == OSTypes.Unix_like:
                path = '/etc/daft/keys'
                fullpath = path + '/{}{}'.format("id_rsa", ext)
            results['results'].append(self.create_directory([s], path))
            results['results'].append(self.push_file(s, file, fullpath))
            res.append(results)
        return res

    def clone_repo(self, server_list, repo_uri, dest):
        cwd, dest_dir = os.path.split(dest)
        return self.sc.salt_command(server_list, 'cmd.run',
                                    ['git clone {uri} {dest}'.format(uri=repo_uri, dest=dest_dir)],
                                    opts={
                                        'cwd': cwd,
                                        'expr_form': 'list'
                                    }, timeout=1200)

    def checkout_branch(self, server_list, location, branch_name):
        return self.sc.salt_command(server_list, 'cmd.run', ['git checkout {}'.format(branch_name)], opts={
            'cwd': location,
            'expr_form': 'list'
        })

    def highstate(self, server_list):
        return self.sc.salt_command(server_list, 'state.highstate', [], opts={'expr_form': 'list'}, timeout=300)

class SaltController:
    def __init__(self, settings):
        self.settings = settings
        self.salt_client = salt.client.LocalClient()
        self.file_roots = None
        self.pillar_roots = None
        self.sls_dir = None

        self.load_salt_info()

    def salt_command(self, target, cmd, arg, opts={}, timeout=120):
        return self.salt_client.cmd(target, cmd, arg, kwarg=opts, timeout=timeout)

    def run_command(self, target, cmd, shell, timeout=120):
        return self.salt_command(target, 'cmd.run', {"arg": [cmd], "shell": shell}, timeout=timeout)

    def load_salt_info(self):
        util.debugLog(self, "load_salt_info")
        master_config = salt.config.master_config(os.environ.get('SALT_MASTER_CONFIG', '/etc/salt/master'))
        self.file_roots = master_config['file_roots']
        self.pillar_roots = master_config['pillar_roots']
        self.sls_dir = "{}/{}".format(self.file_roots['base'][0], self.settings['daft.salt.slsdir'])
        if not os.path.isdir(self.sls_dir):
            os.mkdir(self.sls_dir)
        self.include_daft_top()

    def include_daft_top(self):
        util.debugLog(self, "include_daft_top")
        top_sls = "{}/top.sls".format(self.file_roots['base'][0])
        if not os.path.isfile(top_sls):
            util.debugLog(self, "ERROR: default salt top.sls not found! (check base file roots setting)")
            raise FatalSaltError
        lock = lockfile.FileLock(top_sls)
        lock.acquire(timeout=60)
        with open(top_sls, 'r') as f:
            top_content = yaml.load(f, Loader=Loader)<commit_msg>salt fix<commit_after>import salt.client
import salt.config
import lockfile
import os
import os.path
import random
import shutil
import string
import simplejson as json
import yaml
try:
    from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper

import util

__author__ = 'bkeroack'

class FatalSaltError(Exception):
    pass

class OSTypes:
    Windows = 0
    Unix_like = 1

class RemoteCommands:
    def __init__(self, sc):
        self.sc = sc

    def get_os(self, server):
        resp = self.sc.salt_command(server, 'grains.item', ["os"])
        return OSTypes.Windows if resp[server]['os'] == "Windows" else OSTypes.Unix_like

    def create_directory(self, server_list, path):
        assert isinstance(server_list, list)
        return self.sc.salt_command(server_list, 'file.mkdir', [path])

    def delete_directory_win(self, server_list, path):
        return self.sc.salt_command(server_list, 'cmd.run', ["Remove-Item -Recurse -Force {}".format(path)],
                                    opts={'shell': 'powershell'})

    def delete_directory_unix(self, server_list, path):
        return self.sc.salt_command(server_list, 'cmd.run', ["rm -rf {}".format(path)])

    def delete_directory(self, server_list, path):
        assert isinstance(server_list, list)
        res = list()
        win_s = list()
        unix_s = list()
        for s in server_list:
            ost = self.get_os(s)
            if ost == OSTypes.Windows:
                win_s.append(s)
            elif os == OSTypes.Unix_like:
                unix_s.append(s)
        if len(win_s) > 0:
            res.append(self.delete_directory_win(win_s, path))
        if len(unix_s) > 0:
            res.append(self.delete_directory_unix(unix_s, path))
        return res

    def push_file(self, target, local_path, remote_path):
        r_postfix = ''.join(random.sample(string.letters + string.digits, 20))
        root = self.sc.sls_dir
        newname = os.path.basename(local_path) + r_postfix
        new_fullpath = "{}/{}".format(root, newname)
        shutil.copy(local_path, new_fullpath)
        salt_uri = 'salt://{}/{}'.format(self.sc.settings['daft.salt.slsdir'], newname)
        util.debugLog(self, "push_file: salt_uri: {}".format(salt_uri))
        util.debugLog(self, "push_file: remote_path: {}".format(remote_path))
        res = self.sc.salt_command(target, 'cp.get_file', [salt_uri, remote_path])
        util.debugLog(self, "push_file: resp: {}".format(res))
        os.unlink(new_fullpath)
        return {'success': {'target': target, 'remote_path': remote_path}}

    def push_key(self, server_list, file, name, ext):
        res = list()
        for s in server_list:
            ost = self.get_os(s)
            results = {
                'server': s,
                'os': ost,
                'results': []
            }
            if ost == OSTypes.Windows:
                path = 'C:\Program Files (x86)\Git\.ssh'
                fullpath = path + "\{}{}".format("id_rsa", ext)
            elif ost == OSTypes.Unix_like:
                path = '/etc/daft/keys'
                fullpath = path + '/{}{}'.format("id_rsa", ext)
            results['results'].append(self.create_directory([s], path))
            results['results'].append(self.push_file(s, file, fullpath))
            res.append(results)
        return res

    def clone_repo(self, server_list, repo_uri, dest):
        cwd, dest_dir = os.path.split(dest)
        return self.sc.salt_command(server_list, 'cmd.run',
                                    ['git clone {uri} {dest}'.format(uri=repo_uri, dest=dest_dir)],
                                    opts={'cwd': cwd}, timeout=1200)

    def checkout_branch(self, server_list, location, branch_name):
        return self.sc.salt_command(server_list, 'cmd.run', ['git checkout {}'.format(branch_name)],
                                    opts={'cwd': location})

    def highstate(self, server_list):
        return self.sc.salt_command(server_list, 'state.highstate', [], timeout=300)

class SaltController:
    def __init__(self, settings):
        self.settings = settings
        self.salt_client = salt.client.LocalClient()
        self.file_roots = None
        self.pillar_roots = None
        self.sls_dir = None

        self.load_salt_info()

    def salt_command(self, target, cmd, arg, opts={}, timeout=120):
        return self.salt_client.cmd(target, cmd, arg, kwarg=opts, timeout=timeout, expr_form='list')

    def run_command(self, target, cmd, shell, timeout=120):
        return self.salt_command(target, 'cmd.run', [cmd], opts={"shell": shell}, timeout=timeout)

    def load_salt_info(self):
        util.debugLog(self, "load_salt_info")
        master_config = salt.config.master_config(os.environ.get('SALT_MASTER_CONFIG', '/etc/salt/master'))
        self.file_roots = master_config['file_roots']
        self.pillar_roots = master_config['pillar_roots']
        self.sls_dir = "{}/{}".format(self.file_roots['base'][0], self.settings['daft.salt.slsdir'])
        if not os.path.isdir(self.sls_dir):
            os.mkdir(self.sls_dir)
        self.include_daft_top()

    def include_daft_top(self):
        util.debugLog(self, "include_daft_top")
        top_sls = "{}/top.sls".format(self.file_roots['base'][0])
        if not os.path.isfile(top_sls):
            util.debugLog(self, "ERROR: default salt top.sls not found! (check base file roots setting)")
            raise FatalSaltError
        lock = lockfile.FileLock(top_sls)
        lock.acquire(timeout=60)
        with open(top_sls, 'r') as f:
            top_content = yaml.load(f, Loader=Loader)
{'fraction_non_alphanumeric': 0.13221649484536083, 'fraction_numerical': 0.00429553264604811, 'mean_word_length': 2.4481635071090047, 'pattern_counts': {'":': 3, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1263567', 'n_tokens_mistral': 3822, 'n_tokens_neox': 3661, 'n_words': 764}