<commit_before>'''
Utility functions for saltcloud
'''

# Import python libs
import os
import sys
import shutil
import socket
import tempfile
import time
import subprocess
import salt.utils.event
import multiprocessing
import logging
import types
import re

# Get logging started
log = logging.getLogger(__name__)

# Import salt libs
import salt.crypt
import salt.client
from salt.exceptions import SaltException

# Import third party libs
from jinja2 import Template
import yaml

NSTATES = {
    0: 'running',
    1: 'rebooting',
    2: 'terminated',
    3: 'pending',
}


def os_script(os_, vm_=None, opts=None, minion=''):
    '''
    Return the script as a string for the specific os
    '''
    deploy_path = os.path.join(
        os.path.dirname(os.path.dirname(__file__)),
        'deploy'
    )
    for fn_ in os.listdir(deploy_path):
        full = os.path.join(deploy_path, fn_)
        if not os.path.isfile(full):
            continue
        if os_.lower() == fn_.split('.')[0].lower():
            # found the right script to embed, go for it
            try:
                with open(full, 'r') as fp_:
                    template = Template(fp_.read())
                return str(template.render(opts=opts, vm=vm_, minion=minion))
            except AttributeError:
                # Specified renderer was not found
                continue
    # No deploy script was found, return an empy string
    return ''


def gen_keys(keysize=2048):
    '''
    Generate Salt minion keys and return them as PEM file strings
    '''
    # Mandate that keys are at least 2048 in size
    if keysize < 2048:
        keysize = 2048
    tdir = tempfile.mkdtemp()
    salt.crypt.gen_keys(tdir, 'minion', keysize)
    priv_path = os.path.join(tdir, 'minion.pem')
    pub_path = os.path.join(tdir, 'minion.pub')
    with open(priv_path) as fp_:
        priv = fp_.read()
    with open(pub_path) as fp_:
        pub = fp_.read()
    shutil.rmtree(tdir)
    return priv, pub


def accept_key(pki_dir, pub, id_):
    '''
    If the master config was available then we will have a pki_dir key in
    the opts directory, this method places the pub key in the accepted
    keys dir and removes it from the unaccepted keys dir if that is the case.
    '''
    for key_dir in ('minions', 'minions_pre', 'minions_rejected'):
        key_path = os.path.join(pki_dir, key_dir)
        if not os.path.exists(key_path):
            os.makedirs(key_path)

    key = os.path.join(
        pki_dir,
        'minions/{0}'.format(id_)
    )
    with open(key, 'w+') as fp_:
        fp_.write(pub)

    oldkey = os.path.join(
        pki_dir,
        'minions_pre/{0}'.format(id_)
    )
    if os.path.isfile(oldkey):
        with open(oldkey) as fp_:
            if fp_.read() == pub:
                os.remove(oldkey)


def remove_key(pki_dir, id_):
    '''
    This method removes a specified key from the accepted keys dir
    '''
    key = os.path.join(
        pki_dir,
        'minions/{0}'.format(id_)
    )
    if os.path.isfile(key):
        os.remove(key)


def rename_key(pki_dir, id_, new_id):
    '''
    Rename a key, when an instance has also been renamed
    '''
    oldkey = os.path.join(pki_dir, 'minions/{0}'.format(id_))
    newkey = os.path.join(pki_dir, 'minions/{0}'.format(new_id))
    if os.path.isfile(oldkey):
        os.rename(oldkey, newkey)


def get_option(option, opts, vm_):
    '''
    Convenience function to return the dominant option to be used. Always
    default to options set in the VM structure, but if the option is not
    present there look for it in the main config file
    '''
    if option in vm_:
        return vm_[option]
    if option in opts:
        return opts[option]


def minion_conf_string(opts, vm_):
    '''
    Return a string to be passed into the deployment script for the minion
    configuration file
    '''
    minion = {'id': vm_['name']}
    if 'master_finger' in vm_:
        minion['master_finger'] = vm_['master_finger']
    minion.update(opts.get('minion', {}))
    minion.update(vm_.get('minion', {}))
    if 'master' not in minion:
        raise ValueError("A master was not defined.")
    minion.update(opts.get('map_minion', {}))
    minion.update(vm_.get('map_minion', {}))
    optsgrains = opts.get('map_grains', {})
    if optsgrains:
        minion.setdefault('grains', {}).update(optsgrains)
    vmgrains = vm_.get('map_grains', {})
    if vmgrains:
        minion.setdefault('grains', {}).update(vmgrains)
    return yaml.safe_dump(minion, default_flow_style=False)


def master_conf_string(opts, vm_):
    '''
    Return a string to be passed into the deployment script for the master
    configuration file
    '''
    master = {}

    master.update(opts.get('master', {}))
    master.update(vm_.get('master', {}))

    master.update(opts.get('map_master', {}))
    master.update(vm_.get('map_master', {}))

    return yaml.safe_dump(master, default_flow_style=False)


def wait_for_ssh(host, port=22, timeout=900):
    '''
    Wait until an ssh connection can be made on a specified host
    '''
    start = time.time()
    log.debug('Attempting SSH connection to host {0} on port {1}'.format(host, port))
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    trycount = 0
    while True:
        trycount += 1
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((host, port))
            sock.shutdown(2)
            return True
        except Exception as e:
            log.debug('Caught exception in wait_for_ssh: {0}'.format(e))
            time.sleep(1)
            if time.time() - start > timeout:
                log.error('SSH connection timed out: {0}'.format(timeout))<commit_msg>Remove original/extraneous socket declaration<commit_after>'''
Utility functions for saltcloud
'''

# Import python libs
import os
import sys
import shutil
import socket
import tempfile
import time
import subprocess
import salt.utils.event
import multiprocessing
import logging
import types
import re

# Get logging started
log = logging.getLogger(__name__)

# Import salt libs
import salt.crypt
import salt.client
from salt.exceptions import SaltException

# Import third party libs
from jinja2 import Template
import yaml

NSTATES = {
    0: 'running',
    1: 'rebooting',
    2: 'terminated',
    3: 'pending',
}


def os_script(os_, vm_=None, opts=None, minion=''):
    '''
    Return the script as a string for the specific os
    '''
    deploy_path = os.path.join(
        os.path.dirname(os.path.dirname(__file__)),
        'deploy'
    )
    for fn_ in os.listdir(deploy_path):
        full = os.path.join(deploy_path, fn_)
        if not os.path.isfile(full):
            continue
        if os_.lower() == fn_.split('.')[0].lower():
            # found the right script to embed, go for it
            try:
                with open(full, 'r') as fp_:
                    template = Template(fp_.read())
                return str(template.render(opts=opts, vm=vm_, minion=minion))
            except AttributeError:
                # Specified renderer was not found
                continue
    # No deploy script was found, return an empy string
    return ''


def gen_keys(keysize=2048):
    '''
    Generate Salt minion keys and return them as PEM file strings
    '''
    # Mandate that keys are at least 2048 in size
    if keysize < 2048:
        keysize = 2048
    tdir = tempfile.mkdtemp()
    salt.crypt.gen_keys(tdir, 'minion', keysize)
    priv_path = os.path.join(tdir, 'minion.pem')
    pub_path = os.path.join(tdir, 'minion.pub')
    with open(priv_path) as fp_:
        priv = fp_.read()
    with open(pub_path) as fp_:
        pub = fp_.read()
    shutil.rmtree(tdir)
    return priv, pub


def accept_key(pki_dir, pub, id_):
    '''
    If the master config was available then we will have a pki_dir key in
    the opts directory, this method places the pub key in the accepted
    keys dir and removes it from the unaccepted keys dir if that is the case.
    '''
    for key_dir in ('minions', 'minions_pre', 'minions_rejected'):
        key_path = os.path.join(pki_dir, key_dir)
        if not os.path.exists(key_path):
            os.makedirs(key_path)

    key = os.path.join(
        pki_dir,
        'minions/{0}'.format(id_)
    )
    with open(key, 'w+') as fp_:
        fp_.write(pub)

    oldkey = os.path.join(
        pki_dir,
        'minions_pre/{0}'.format(id_)
    )
    if os.path.isfile(oldkey):
        with open(oldkey) as fp_:
            if fp_.read() == pub:
                os.remove(oldkey)


def remove_key(pki_dir, id_):
    '''
    This method removes a specified key from the accepted keys dir
    '''
    key = os.path.join(
        pki_dir,
        'minions/{0}'.format(id_)
    )
    if os.path.isfile(key):
        os.remove(key)


def rename_key(pki_dir, id_, new_id):
    '''
    Rename a key, when an instance has also been renamed
    '''
    oldkey = os.path.join(pki_dir, 'minions/{0}'.format(id_))
    newkey = os.path.join(pki_dir, 'minions/{0}'.format(new_id))
    if os.path.isfile(oldkey):
        os.rename(oldkey, newkey)


def get_option(option, opts, vm_):
    '''
    Convenience function to return the dominant option to be used. Always
    default to options set in the VM structure, but if the option is not
    present there look for it in the main config file
    '''
    if option in vm_:
        return vm_[option]
    if option in opts:
        return opts[option]


def minion_conf_string(opts, vm_):
    '''
    Return a string to be passed into the deployment script for the minion
    configuration file
    '''
    minion = {'id': vm_['name']}
    if 'master_finger' in vm_:
        minion['master_finger'] = vm_['master_finger']
    minion.update(opts.get('minion', {}))
    minion.update(vm_.get('minion', {}))
    if 'master' not in minion:
        raise ValueError("A master was not defined.")
    minion.update(opts.get('map_minion', {}))
    minion.update(vm_.get('map_minion', {}))
    optsgrains = opts.get('map_grains', {})
    if optsgrains:
        minion.setdefault('grains', {}).update(optsgrains)
    vmgrains = vm_.get('map_grains', {})
    if vmgrains:
        minion.setdefault('grains', {}).update(vmgrains)
    return yaml.safe_dump(minion, default_flow_style=False)


def master_conf_string(opts, vm_):
    '''
    Return a string to be passed into the deployment script for the master
    configuration file
    '''
    master = {}

    master.update(opts.get('master', {}))
    master.update(vm_.get('master', {}))

    master.update(opts.get('map_master', {}))
    master.update(vm_.get('map_master', {}))

    return yaml.safe_dump(master, default_flow_style=False)


def wait_for_ssh(host, port=22, timeout=900):
    '''
    Wait until an ssh connection can be made on a specified host
    '''
    start = time.time()
    log.debug('Attempting SSH connection to host {0} on port {1}'.format(host, port))
    trycount = 0
    while True:
        trycount += 1
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((host, port))
            sock.shutdown(2)
            return True
        except Exception as e:
            log.debug('Caught exception in wait_for_ssh: {0}'.format(e))
            time.sleep(1)
            if time.time() - start > timeout:
                log.error('SSH connection timed out: {0}'.format(timeout))
{'fraction_non_alphanumeric': 0.11757501530924679, 'fraction_numerical': 0.0069985128160265945, 'mean_word_length': 2.5725, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '786910', 'n_tokens_mistral': 3761, 'n_tokens_neox': 3640, 'n_words': 1037}