#!python
from pyngen import *
from argparse import *
import os
import sys
import click


parser = ArgumentParser()
parser.add_argument('revs', metavar='function', choices=[
                    "newincident", "init"], help='choice function')
parser.add_argument("-feed", dest="feed", help="Select feed slug",
                    required='newincident' in sys.argv)
parser.add_argument("-address", dest="address", help="IP or Domain",
                    required='newincident' in sys.argv)
parser.add_argument("-type", dest="type", help="Select type slug",
                    required='newincident' in sys.argv)

evidence_group = parser.add_mutually_exclusive_group()
evidence_group.add_argument("-evidenceFile", dest="evidenceFile",
                            help="Select path of evidence")
evidence_group.add_argument("-evidenceText", dest="evidenceText",
                            help="Select text to send as evidence", default=None)

parser.add_argument("-apikey", dest="apikey",
                    help="Set API KEY for init", required='init' in sys.argv)
parser.add_argument("-url", dest="url",
                    help="Set Ngen URL for init", required='init' in sys.argv)
args = parser.parse_args()


NGEN_CONFIG_DIR = '~/.ngen/'
COLORIZE_FIELDS = {
    'ip_str': 'green',
    'port': 'yellow',
    'data': 'white',
    'hostnames': 'magenta',
    'org': 'cyan',
    'vulns': 'red',
}


ngen_dir = os.path.expanduser(NGEN_CONFIG_DIR)


def inits(url, apikey):
    """Initialize the Ngen command-line"""
    # Create the directory if necessary
    if not os.path.isdir(ngen_dir):
        try:
            os.mkdir(ngen_dir)
        except OSError:
            raise click.ClickException(
                'Unable to create directory to store the Ngen API key ({})'.format(ngen_dir))

    # Make sure it's a valid API key
    url = url.strip()
    apikey = apikey.strip()
    try:
        api = PyNgen(url, apikey)
        # api.info()
    except pyngen.UnauthorizedNgenError as e:
        raise click.ClickException(e.detail)

    # Store the API key in the user's directory
    keyfile = ngen_dir + '/apikey'
    with open(keyfile, 'w') as fout:
        fout.write(apikey)
    os.chmod(keyfile, 0o600)

    keyfile = ngen_dir + '/url'
    with open(keyfile, 'w') as fout:
        fout.write(url)
    os.chmod(keyfile, 0o600)


def createNewIncident(pyngen, address, incident_feed, incident_type, **kargs):
    try:
        res = pyngen.newIncident(
            address, incident_feed, incident_type, **kargs)
        print("Added incident ID: {}".format(res))
    except NewIncidentFieldError as e:
        print("Error creating incident: \n{}".format(e.detail))
        exit(2)


if (args.revs == "newincident"):
    try:
        apikey = open(ngen_dir + '/apikey').read()
        url = open(ngen_dir + '/url').read()
    except:
        print("Please use init first")
        exit(2)

    ngen = PyNgen(url, apikey)
    createNewIncident(ngen, args.address, args.feed,
                      args.type, evidence_text=args.evidenceText, evidence_file=args.evidenceFile)
elif (args.revs == "init"):
    inits(args.url, args.apikey)
