#!/usr/bin/env python
# Uploads reports from the Salt job cache to Foreman

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

import httplib
import ssl
import json
import yaml
import os
import sys
import base64

import traceback

import salt.config
import salt.runner


def salt_config():
    with 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.iteritems():
        if type(value) == str:
            result[minion] = [value]

    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 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 open(LAST_UPLOADED, 'w+') as f:
        f.write(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 open(os.devnull, 'wb') 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",
    })
    last_uploaded = read_last_uploaded()

    job_ids = [jid for (jid, value) in jobs.iteritems()
               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 = httplib.HTTPSConnection(config[':host'],
                port=config[':port'], context=ctx)
    else:
        connection = httplib.HTTPConnection(config[':host'],
                port=config[':port'])
        if ':username' in config and ':password' in config:
            token = base64.b64encode('{}:{}'.format(config[':username'],
                                                    config[':password']))
            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:
        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()
