#!/usr/bin/env python

"""CLI for submitting tasks to Retrace Server."""

from __future__ import print_function
import os
import sys
import time
import shutil
import socket
import tarfile
import logging
import argparse
import tempfile
import requests
from requests.packages.urllib3.exceptions import InsecureRequestWarning
from requests_kerberos import HTTPKerberosAuth, OPTIONAL


TASK_RETRACE, TASK_DEBUG, TASK_VMCORE, TASK_RETRACE_INTERACTIVE, \
        TASK_VMCORE_INTERACTIVE = range(5)


def ticket_check():
    """Check if the user has a valid kerberose ticket."""
    return not os.system("klist -s")


def prepare_tar(args):
    """Create tar file containg all necessary files."""
    LOGGER.info("Preparing tar")

    tmpdir = tempfile.mkdtemp()
    tar_name = os.path.join(tmpdir, "archive.tar.gz")
    tar_file = tarfile.open(tar_name, "w:gz")
    for itm in ["vra", "os_release", "package", "executable"]:
        if itm in args.keys() and args[itm]:
            LOGGER.info("Adding '%s' into tar", format(itm))
            fl_name = os.path.join(tmpdir, itm)
            with open(fl_name, "w") as fld:
                fld.write(args[itm])
            tar_file.add(fl_name, arcname=itm)

    LOGGER.info("Adding '%s' as '%s' into tar",
                args['COREFILE'],
                args['task_type'])
    if not os.path.isfile(args['COREFILE']):
        print("'{0}' does not exists".format(args['COREFILE']))
        sys.exit(1)
    tar_file.add(args['COREFILE'], arcname=args['task_type'])
    tar_file.close()
    LOGGER.info("Tar file created as '%s'", tar_name)
    return tar_name


def check_response(response, kill=True):
    """Test return code of response.

    If @kill is True. then when code differ from 200 exit the whole script.
    """
    if response.status_code not in [200, 201]:
        LOGGER.info("Code %s was returned. Headers:\n%s\nText:\n%s\n",
                    response.status_code, response.headers, response.text)
        if kill:
            print("Status code {0} was returned with message: '{1}'"
                  .format(response.status_code, response.text))
            sys.exit(1)
        else:
            return 1
    return 0


def parse_task_id(response):
    """Parse Task_ID from manager page."""
    for line in response.text.split("\n"):
        if "<title>Task #" in line:
            return line.split()[1][1:]
    return "unknown"


def parse_status(response):
    """Parse status from manager page."""
    lines = response.text.split("\n")
    for i, line in enumerate(lines):
        if "<th>Status:</th>" in line:
            return lines[i+1].strip()[4:-5]
    return "unknown"


def parse_md5sum(response):
    """Parse md5sum from manager page."""
    lines = response.text.split("\n")
    for line in lines:
        if "<th>Md5sum:</th>" in line:
            return line[line.find("<td>")+4:]
    return "Md5sum not calculated. Not yet downloaded or md5sum not requested."


def parse_email(response):
    """Parse email from manager page."""
    key_string = 'Update e-mail(s)'
    lines = response.text.split("\n")
    for line in lines:
        if key_string in line:
            first = line.find("value=")+7
            last = line.find('"', first)
            candidate = line[first:last]
            if candidate == key_string:
                return "Email is not set."
            return candidate


def parse_caseno(response):
    """Parse caseno from manager page."""
    key_string = 'Update case no.'
    lines = response.text.split("\n")
    for line in lines:
        if key_string in line:
            first = line.find("value=")+7
            last = line.find('"', first)
            candidate = line[first:last]
            if candidate == key_string:
                return "Case number is not set."
            return candidate

def parse_bugzillano(response):
    """Parse bugzillano from manager page."""
    key_string = 'Update bugzilla no.'
    lines = response.text.split("\n")
    for line in lines:
        if key_string in line:
            first = line.find("value=")+7
            last = line.find('"', first)
            candidate = line[first:last]
            if candidate == key_string:
                return "Bugzilla number is not set."
            return candidate


def create_new_task(args):
    """Start new task."""
    if args['task_input'] == "http":
        LOGGER.info("Starting new http task")
        if args['task_type'] == 'coredump':
            if args['interactive']:
                task_type = TASK_RETRACE_INTERACTIVE
            else:
                task_type = TASK_RETRACE
        elif args['interactive']:
            task_type = TASK_VMCORE_INTERACTIVE
        else:
            task_type = TASK_VMCORE
        LOGGER.info("Task_type: %s", task_type)

        tar_name = prepare_tar(args)
        tar_size = os.path.getsize(tar_name)
        LOGGER.info("Tar size: %s", tar_size)

        headers = {'Content-Type': 'application/x-tar',
                   'Content-Length': tar_size,
                   'X-Task-Type': task_type,
                   'Connection': 'close'}

        req = requests.Request('POST',
                               args['server']+"/create",
                               headers=headers)

        LOGGER.info("Connecting to %s", args['server']+"/create")

        prepped = req.prepare()
        with open(tar_name, 'rb') as fld:
            prepped.body = fld.read(tar_size)

        session = requests.Session()

        res = session.send(prepped,
                           stream=True,
                           verify=(not args['no_verify']))

        shutil.rmtree(os.path.dirname(tar_name))

        check_response(res)

        task_id = res.headers['X-Task-ID']
        task_password = res.headers['X-Task-Password']

        print("Task created")
        print("Task ID: {0}".format(task_id))
        print("Task password: {0}".format(task_password))
        return (task_id, task_password)

    elif args['task_input'] == "manager":
        args['COREFILE'] = "file://" + os.path.abspath(args['COREFILE'])
        LOGGER.info("Starting new manager task")
        server = args['server'] + "/manager/"
        payload = {"custom_url": args['COREFILE']}
        payload['md5sum'] = 'off' if args['no_md5'] else 'on'
        for itm in ["vra", "task_type", "os_release", "package", "executable"]:
            if itm in args.keys():
                payload[itm] = args[itm]

        LOGGER.info("Connecting to %s", server + "__custom__")

        res = requests.post(server + "__custom__",
                            data=payload,
                            verify=(not args['no_verify']),
                            auth=args['kerberos'])

    elif args['task_input'] == "ftp":
        LOGGER.info("Starting new ftp task")
        server = args['server'] + "/manager/"
        payload = {'md5sum': 'off' if args['no_md5'] else 'on'}
        for itm in ["vra", "task_type", "os_release", "package", "executable"]:
            if itm in args.keys():
                payload[itm] = args[itm]

        LOGGER.info("Connecting to server %s",
                    server + args['COREFILE'] + "/start")
        res = requests.get(server + args['COREFILE'] + "/start",
                           params=payload,
                           verify=(not args['no_verify']),
                           auth=args['kerberos'])

    check_response(res)
    task_id = parse_task_id(res)
    print("Task created on {0}"
          .format(server + task_id))

    ARGS['TASK_ID'] = task_id
    ARGS['password'] = None

    if ARGS['email']:
        update_email(ARGS)

    if ARGS['caseno']:
        update_caseno(ARGS)

    if ARGS['bugzillano']:
        update_bugzillano(ARGS)

    return (task_id, None)


def get_status(args):
    """Get current task's status."""
    if args['password']:
        LOGGER.info("Getting status from http task on server %s",
                    args['server'] + "/" + args['TASK_ID'])
        server = args['server'] + "/" + args['TASK_ID']
        headers = {"X-Task-Password": args['password']}
        res = requests.get(server,
                           headers=headers,
                           verify=(not args['no_verify']),
                           auth=args['kerberos'])

        check_response(res)
        status = res.text
        return status

    LOGGER.info("Getting status from manager/ftp task on server %s",
                args['server'] + "/manager/" + args['TASK_ID'])
    server = args['server'] + "/manager/" + args['TASK_ID']
    res = requests.get(server,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    check_response(res)
    status = parse_status(res)
    return status


def get_md5sum(args):
    """Get md5sum of uploaded file."""
    # Only supported for manager tasks
    if args['password']:
        print("Md5sum are supported for manager and ftp tasks")
        sys.exit(1)

    LOGGER.info("Getting md5sum from server %s",
                args['server'] + "/manager/" + args['TASK_ID'])

    server = args['server'] + "/manager/" + args['TASK_ID']
    res = requests.get(server,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    if check_response(res, False):
        print("Error: md5sum is only supported for manager and ftp tasks")
        return 0
    md5sum = parse_md5sum(res)
    print(md5sum)
    return 1


def get_email(args):
    """Get email of a task."""
    # Only supported for manager tasks
    if args['password']:
        print("Emails are supported for manager and ftp tasks")
        sys.exit(1)

    LOGGER.info("Getting email from server %s",
                args['server'] + "/manager/" + args['TASK_ID'])

    server = args['server'] + "/manager/" + args['TASK_ID']
    res = requests.get(server,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    if check_response(res, False):
        print("Error: email is only supported for manager and ftp tasks")
        return 0
    email = parse_email(res)
    print(email)
    return 1


def get_caseno(args):
    """Get case number of a task."""
    # Only supported for manager tasks
    if args['password']:
        print("Case numbers are supported for manager and ftp tasks")
        sys.exit(1)

    LOGGER.info("Getting case number from server %s",
                args['server'] + "/manager/" + args['TASK_ID'])

    server = args['server'] + "/manager/" + args['TASK_ID']
    res = requests.get(server,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    if check_response(res, False):
        print("Error: caseno is only supported for manager and ftp tasks")
        return 0
    caseno = parse_caseno(res)
    print(caseno)
    return 1

def get_bugzillano(args):
    """Get bugzilla number of a task."""
    # Only supported for manager tasks
    if args['password']:
        print("Bugzilla numbers are supported for manager and ftp tasks")
        sys.exit(1)

    LOGGER.info("Getting bugzilla number from server %s",
                args['server'] + "/manager/" + args['TASK_ID'])

    server = args['server'] + "/manager/" + args['TASK_ID']
    res = requests.get(server,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    if check_response(res, False):
        print("Error: bugzillano is only supported for manager and ftp tasks")
        return 0
    bugzillano = parse_bugzillano(res)
    print(bugzillano)
    return 1

def get_backtrace(args):
    """Get a backtrace of the task."""
    if args['password']:
        server = args['server'] + "/" + args['TASK_ID'] + "/backtrace"
        headers = {"X-Task-Password": args['password']}

    else:
        server = args['server'] + "/manager/" + args['TASK_ID'] + "/backtrace"
        headers = {}

    LOGGER.info("Getting backtrace from server %s", server)

    res = requests.get(server,
                       headers=headers,
                       verify=(not args['no_verify']),
                       auth=args['kerberos'])

    if check_response(res, False):
        if res.text == "There is no backtrace for the specified task":
            print("No backtrace. Check status if task is running or failed.")
            return 0
        else:
            print("Error: {0}".format(res.text))
            exit(1)

    with open("backtrace", "w") as fld:
        fld.write(res.text)
    print("Backtrace saved into 'backtrace' file")
    return 1


def run_batch(args):
    """Run the whole retrace process in one go."""
    task_id, password = create_new_task(args)
    args['password'] = password
    args['TASK_ID'] = task_id
    last_status = ""
    new_status = get_status(args)
    while True:
        print(new_status)
        last_status = new_status
        if new_status == "Retrace job finished successfully":
            break
        if new_status == "Retrace job failed":
            sys.exit(1)
        while new_status == last_status:
            time.sleep(10)
            new_status = get_status(args)

    get_backtrace(args)


def update_email(args):
    """Update the email address of the task."""
    # Only supported for manager tasks
    if args['password']:
        print("Email notification only supported for manager and ftp tasks")
        sys.exit(1)
    server = args['server'] + "/manager/" + args['TASK_ID'] + "/notify"
    res = requests.post(server,
                        data={"notify": args['email']},
                        verify=(not args['no_verify']),
                        auth=args['kerberos'])

    check_response(res)
    print("Added email address '{0}' to the task '{1}'"
          .format(args['email'], args['TASK_ID']))


def update_caseno(args):
    """Update case number of the task."""
    # Only supported for manager tasks
    if args['password']:
        print("Case number is only supported for manager and ftp tasks")
        sys.exit(1)
    server = args['server'] + "/manager/" + args['TASK_ID'] + "/caseno"
    res = requests.post(server,
                        data={"caseno": args['caseno']},
                        verify=(not args['no_verify']),
                        auth=args['kerberos'])

    check_response(res)
    print("Added case number '{0}' to the task '{1}'"
          .format(args['caseno'], args['TASK_ID']))

def update_bugzillano(args):
    """Update bugzilla number of the task."""
    # Only supported for manager tasks
    if args['password']:
        print("Case number is only supported for manager and ftp tasks")
        sys.exit(1)
    server = args['server'] + "/manager/" + args['TASK_ID'] + "/bugzillano"
    res = requests.post(server,
                        data={"bugzillano": args['bugzillano']},
                        verify=(not args['no_verify']),
                        auth=args['kerberos'])

    check_response(res)
    print("Added bugzilla number '{0}' to the task '{1}'"
          .format(args['bugzillano'], args['TASK_ID']))

if __name__ == "__main__":
    # parse arguments
    PARSER = argparse.ArgumentParser(description=("Create, watch and set"
                                                  "tasks through commandline"))
    TASK_PARSER = PARSER.add_subparsers(dest='task_action')
    for cmd in ['create', 'batch']:
        tmp = TASK_PARSER.add_parser(cmd)
        type_group = tmp.add_mutually_exclusive_group()
        type_group.add_argument("-u", "--userspace_core", action="store_const",
                                dest="task_type", const="coredump",
                                help="Retrace userspace coredump")
        type_group.add_argument("-k", "--kernel_vmcore", action="store_const",
                                dest="task_type", const="vmcore",
                                help="Retrace kernel vmcore")
        tmp.add_argument("COREFILE",
                         help="File to be retraced")

        input_group = tmp.add_mutually_exclusive_group()
        input_group.add_argument("-t", "--http", action="store_const",
                                 dest="task_input", const="http",
                                 help="Submit local corefile through http")
        input_group.add_argument("-m", "--manager", action="store_const",
                                 dest="task_input", const="manager",
                                 help=("Submit corefile URL that is wget-able"
                                       "or local server file"))
        input_group.add_argument("-f", "--ftp", action="store_const",
                                 dest="task_input", const="ftp",
                                 help=("Submit corefile that is accessible"
                                       "from FTP set on server"))

        tmp.add_argument("-i", "--interactive", action="store_true",
                         help="Create interactive task")
        tmp.add_argument("-d", "--no-md5", action="store_true",
                         help="Do not calculate md5sum on the archive")
        tmp.add_argument("-s", "--server",
                         help="URL for retrace-server")
        tmp.add_argument("-n", "--no-verify", action="store_true",
                         help="Do not verify certificate")
        tmp.add_argument("-v", "--verbose", action="store_const",
                         default=logging.WARNING, const=logging.DEBUG)

        tmp.add_argument("-r", "--kernelver",
                         help="Version of kernel in vmcore")
        tmp.add_argument("-o", "--os_release",
                         help="Version of operating system")
        tmp.add_argument("-p", "--package",
                         help="Version of crashed package")
        tmp.add_argument("-x", "--executable",
                         help="Crash executable")

        tmp.add_argument("-e", "--email",
                         help=("Set email to the task."
                               "Only takes affect with -f or -m"))
        tmp.add_argument("-c", "--caseno",
                         help=("Set caseno to the task."
                               "Only takes affect with -f or -m"))
        tmp.add_argument("-bz", "--bugzillano",
                         help=("Set bugzillano to the task."
                               "Only takes affect with -f or -m"))

    TMP = TASK_PARSER.add_parser('get')
    TMP.add_argument("TASK_ID",
                     help="ID of existing task")
    TMP.add_argument("-p", "--password",
                     help="Password of the task")
    TMP.add_argument("-s", "--server",
                     help="URL for retrace-server")
    TMP.add_argument("-n", "--no-verify", action="store_true",
                     help="Do not verify certificate")
    TMP.add_argument("-v", "--verbose", action="store_const",
                     default=logging.WARNING, const=logging.DEBUG)
    GET_TYPE = TMP.add_mutually_exclusive_group()
    GET_TYPE.add_argument("-t", "--status", action="store_const",
                          dest="get_type", const="status",
                          help="Get current status of the task.")
    GET_TYPE.add_argument("-b", "--backtrace", action="store_const",
                          dest="get_type", const="backtrace",
                          help="Get backtrace of the task.")
    GET_TYPE.add_argument("-m", "--md5sum", action="store_const",
                          dest="get_type", const="md5sum",
                          help="Get md5sum of the task.")
    GET_TYPE.add_argument("-e", "--email", action="store_const",
                          dest="get_type", const="email",
                          help="Get set email of the task.")
    GET_TYPE.add_argument("-c", "--caseno", action="store_const",
                          dest="get_type", const="caseno",
                          help="Get set case number of the task.")
    GET_TYPE.add_argument("-bz", "--bugzillano", action="store_const",
                          dest="get_type", const="bugzillano",
                          help="Get set bugzilla number of the task.")

    TMP = TASK_PARSER.add_parser('set')
    TMP.add_argument("TASK_ID",
                     help="ID of existing task")
    TMP.add_argument("-p", "--password",
                     help="Password of the task")
    TMP.add_argument("-s", "--server",
                     help="URL for retrace-server")
    TMP.add_argument("-n", "--no-verify", action="store_true",
                     help="Do not verify certificate")
    TMP.add_argument("-v", "--verbose", action="store_const",
                     default=logging.WARNING, const=logging.DEBUG)
    TMP.add_argument("-c", "--caseno",
                     help="Set case number of the task.")
    TMP.add_argument("-e", "--email",
                     help="Set email to the task.")
    TMP.add_argument("-bz", "--bugzillano",
                     help="Set bugzilla number of the task.")

    ARGS = vars(PARSER.parse_args())

    # if not verifying certificates, suppress warning
    if ARGS['no_verify']:
        requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

    # create logger
    logging.basicConfig(level=ARGS['verbose'])
    LOGGER = logging.getLogger()

    # check kerberos
    if ticket_check():
        LOGGER.info("Kerberos ticket is valid.")
        ARGS['kerberos'] = HTTPKerberosAuth(mutual_authentication=OPTIONAL)
    else:
        ARGS['kerberos'] = None

    # Try guessing type of task if no specified
    if ARGS['task_action'] in ['create', 'batch'] and not ARGS['task_input']:
        # If the file does not exist locally
        if not os.path.isfile(ARGS['COREFILE']):
            # and server not know
            if not ARGS['server']:
                LOGGER.info("Defaulting to ftp task")
                ARGS['task_input'] = 'ftp'
            else:
                LOGGER.info("Defaulting to manager task")
                ARGS['task_input'] = 'manager'
        else:
            if not ARGS['server']:
                LOGGER.info("Defaulting to manager task")
                ARGS['task_input'] = 'manager'
            else:
                LOGGER.info("Defaulting to http task")
                ARGS['task_input'] = 'http'

    # default to local server if no assigned
    if not ARGS['server']:
        ARGS['server'] = "https://" + socket.gethostname()
        try:
            if requests.get(ARGS['server'],
                            verify=(not ARGS['no_verify']),
                            auth=ARGS['kerberos']).status_code != 200:
                raise requests.exceptions.ConnectionError
        except requests.exceptions.RequestException:
            LOGGER.info("Local server '%s' is not available", ARGS['server'])
            ARGS['server'] = 'https://retrace.fedoraproject.org'
        LOGGER.info("Trying server '%s'", ARGS['server'])

    # HACK - kernelver is better to use than vra
    ARGS['vra'] = ARGS.get('kernelver')

    if ARGS['task_action'] == "create":
        create_new_task(ARGS)

    if ARGS['task_action'] == "batch":
        run_batch(ARGS)

    if ARGS['task_action'] == "get":
        if ARGS['get_type'] == "md5sum":
            get_md5sum(ARGS)
        elif ARGS['get_type'] == "backtrace":
            get_backtrace(ARGS)
        elif ARGS['get_type'] == "email":
            get_email(ARGS)
        elif ARGS['get_type'] == "caseno":
            get_caseno(ARGS)
        elif ARGS['get_type'] == "bugzillano":
            get_bugzillano(ARGS)
        else:
            print("Status: {0}".format(get_status(ARGS)))

    if ARGS['task_action'] == "set":
        if ARGS['email']:
            update_email(ARGS)
        if ARGS['caseno']:
            update_caseno(ARGS)
        if ARGS['bugzillano']:
            update_bugzillano(ARGS)
