#!/usr/local/bin/python3.9
# Uploads reports from the Salt job cache to Foreman

from __future__ import print_function

LAST_UPLOADED = '/usr/local/etc/salt/last_uploaded'
FOREMAN_CONFIG = '/usr/local/etc/salt/foreman.yaml'
LOCK_FILE = '/var/run/salt-report-upload.lock'

try:
    from http.client import HTTPConnection, HTTPSConnection
except ImportError:
    from httplib import HTTPSConnection, HTTPSConnection
import ssl
import json
import yaml
import io
import os
import sys
import base64
import traceback
import salt.config
import salt.runner

if sys.version_info.major == 3:
    unicode = str


def salt_config():
    with io.open(FOREMAN_CONFIG, 'r') as f:
        config = yaml.load(f.read())
    return config


def get_job(job_id):
    result = run('jobs.lookup_jid', [job_id])
    # If any minion's results are strings, they're exceptions
    # and should be wrapped in a list like other errors

    for minion, value in result.items():
        try:
            if isinstance(value,str):
                result[minion] = [value]
            elif isinstance(value,list):
                result[minion] = value
            else:
                for key, entry in value.items():
                    if key.startswith('module_') and '__id__' in entry and entry['__id__'] == 'state.highstate':
                        result[minion] = entry['changes']['ret']
                        break
        except KeyError:
            traceback.print_exc()

    return {'job':
             {
               'result': result,
               'function': 'state.highstate',
               'job_id': job_id
             }
           }


def read_last_uploaded():
    if not os.path.isfile(LAST_UPLOADED):
        return 0
    else:
        with io.open(LAST_UPLOADED, 'r') as f:
            result = f.read().strip()
        if len(result) == 20:
            try:
                return int(result)
            except ValueError:
                return 0
        else:
            return 0


def write_last_uploaded(last_uploaded):
    with io.open(LAST_UPLOADED, 'w+') as f:
        f.write(unicode(last_uploaded))


def run(*args, **kwargs):
    __opts__ = salt.config.master_config(
            os.environ.get('SALT_MASTER_CONFIG', '/etc/salt/master'))

    runner = salt.runner.Runner(__opts__)
    with io.open(os.devnull, 'w') as f:
        stdout_bak, sys.stdout = sys.stdout, f
        try:
            ret = runner.cmd(*args, **kwargs)
        finally:
            sys.stdout = stdout_bak
    return ret['data'] if 'data' in ret else ret


def jobs_to_upload():
    jobs = run('jobs.list_jobs', kwarg={
        "search_function": ["state.highstate","state.template_str"],
    })
    last_uploaded = read_last_uploaded()

    job_ids = [jid for jid in jobs.keys() if int(jid) > last_uploaded]

    for job_id in sorted(job_ids):
        yield job_id, get_job(job_id)


def upload(jobs):
    config = salt_config()
    headers = {'Accept': 'application/json',
               'Content-Type': 'application/json'}

    if config[':proto'] == 'https':
        ctx = ssl.create_default_context()
        ctx.load_cert_chain(certfile=config[':ssl_cert'], keyfile=config[':ssl_key'])
        if config[':ssl_ca']:
            ctx.load_verify_locations(cafile=config[':ssl_ca'])
        connection = HTTPSConnection(config[':host'],
                                     port=config[':port'], context=ctx)
    else:
        connection = HTTPConnection(config[':host'],
                                    port=config[':port'])
        if ':username' in config and ':password' in config:
            auth = '{}:{}'.format(config[':username'], config[':password'])
            if not isinstance(auth, bytes):
                auth = auth.encode('UTF-8')
            token = base64.b64encode(auth)
            headers['Authorization'] = 'Basic {}'.format(token)

    for job_id, job in jobs:

        if job['job']['result'] == {}:
            continue

        connection.request('POST', '/salt/api/v2/jobs/upload',
                json.dumps(job), headers)
        response = connection.getresponse()

        if response.status == 200:
            write_last_uploaded(job_id)
            print("Success %s: %s" % (job_id, response.read()))
        else:
            print("Unable to upload job - aborting report upload")
            print(response.read())


def get_lock():
    if os.path.isfile(LOCK_FILE):
        raise Exception("Unable to obtain lock.")
    else:
        io.open(LOCK_FILE, 'w+').close()


def release_lock():
    if os.path.isfile(LOCK_FILE):
        os.remove(LOCK_FILE)


if __name__ == '__main__':
    try:
        get_lock()
        upload(jobs_to_upload())
        release_lock()
    except:
        release_lock()
        traceback.print_exc()
