salt-highstate-pp

#!/usr/bin/env python3
# salt-highstate-pp (part of ossobv/vcutil) // wdoekes/2016 // Public Domain
import json
import os
import select
import subprocess
import sys
import yaml


class StdioStderrProcessor(object):
    def __init__(self, stdout, stderr):
        self.stdout = stdout.fileno()
        self.stderr = stderr.fileno()
        self.poller = select.poll()
        self.poller.register(stdout)
        self.poller.register(stderr)

    def poll(self, wait):
        res = self.poller.poll(wait)
        res.sort(key=(lambda x: x[0] != self.stderr))  # stderr first
        return res

    def process(self):
        while self.stdout or self.stderr:
            res = self.poll(1000)
            for fileno, value in res:
                if value & select.POLLIN:
                    data = os.read(fileno, 8192)
                    data = data.decode('utf-8', 'replace')
                    self.add_data(fileno, data)
                    break
                if value & select.POLLHUP:
                    self.poller.unregister(fileno)
                    if fileno == self.stdout:
                        self.stdout = None
                    else:
                        assert fileno == self.stderr
                        self.stderr = None
                    break

    def add_data(self, fileno, data):
        if fileno == self.stdout:
            self.add_stdout(data)
        else:
            assert fileno == self.stderr
            self.add_stderr(data)

    def add_stdout(self, data):
        sys.stdout.write(data)

    def add_stderr(self, data):
        sys.stderr.write('\x1b[1;31m{}\x1b[0m'.format(data))
        sys.stderr.flush()

    def finalize(self, retcode):
        if retcode != 0:
            self.add_stderr('[return value {}]'.format(retcode))


class JsonProcessor(StdioStderrProcessor):
    def __init__(self, *args, **kwargs):
        super(JsonProcessor, self).__init__(*args, **kwargs)
        self.jsonbuf = []
        self.success = {}
        self.failure = {}

    def add_stdout(self, data):
        for line_idx, line in enumerate(data.split('\n')):
            if line_idx:
                self.jsonbuf.append('\n')

            self.jsonbuf.append(line)

            if line == '}':
                jsonbuf = ''.join(self.jsonbuf)
                try:
                    dict_ = json.loads(jsonbuf)
                except ValueError:
                    print(jsonbuf)
                    raise
                self.add_server(dict_)
                self.jsonbuf = []

    def add_server(self, dict_):
        keys = tuple(dict_.keys())
        assert len(keys) == 1, dict_
        hostname = keys[0]
        try:
            list_of_results = dict_[hostname].values()
        except AttributeError:
            success = False
        else:
            if all(i['result'] for i in list_of_results):
                success = True
            else:
                success = False

        if success:
            self.add_success(hostname, dict_)
        else:
            self.add_failure(hostname, dict_)

    def add_success(self, hostname, dict_):
        self.success[hostname] = dict_
        print('[\x1b[1;32m OK \x1b[0m] {}'.format(hostname))

    def add_failure(self, hostname, dict_):
        self.failure[hostname] = dict_
        print('[\x1b[1;31mFAIL\x1b[0m] {}'.format(hostname))

    def finalize(self, retcode):
        assert retcode == 0, retcode

        if self.failure:
            for hostname in sorted(self.failure.keys()):
                failfile = '/tmp/{}.saltfail'.format(hostname)
                data = yaml.dump(self.failure[hostname])
                print()
                print('---- {} ({} failure log) ----'.format(
                    failfile, hostname))
                print()
                print('\x1b[1;31m{}\x1b[0m'.format(data))
                try:
                    with open(failfile, 'w') as fp:
                        fp.write(data)
                except:
                    print('(failed writing to {})'.format(failfile))
            retcode = 1

            print('---- summary ----')
            print('\x1b[1m{} success\x1b[0m'.format(len(self.success)))
            print('\x1b[1m{} failed\x1b[0m'.format(len(self.failure)))

        if any(i.strip() for i in self.jsonbuf):
            print()
            print('---- trailing output ----')
            print()
            print(''.join(self.jsonbuf))
            retcode = 2

        return retcode


def salt_highstate(target, *arguments):
    function = 'state.highstate'
    with open('/dev/null', 'r') as NULLI:
        with subprocess.Popen(
                ['salt', '--out=json', target, function] + list(arguments),
                stdin=NULLI, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                env={'PYTHONUNBUFFERED': '1'}) as proc:
            processor = JsonProcessor(stdout=proc.stdout, stderr=proc.stderr)
            processor.process()
            retcode = proc.wait()
            retcode = processor.finalize(retcode)
    sys.exit(retcode)


if __name__ == '__main__':
    salt_highstate(*sys.argv[1:])