#!/usr/bin/python3
from pyngen import *
from argparse import *
import os
import sys
import click


parser = ArgumentParser(formatter_class=RawTextHelpFormatter, epilog="""
EXAMPLES:
%(prog)s init -apikey 0123456789abcdef -url https://ngen.example.com/api
%(prog)s newincident -feed external_report -type malware -address ip.or.domain.com
%(prog)s reportcsv -feed external_report -type miner -path path/to/file.csv  -addressheader ip""")

parser.add_argument('revs', metavar='action', choices=[
                    "newincident", "init", "reportcsv"], help='Set action: %(choices)s')

newincident_group = parser.add_argument_group('newincident arguments')

feed = newincident_group.add_argument("-feed", dest="feed", help="Select feed slug",
required='newincident' in sys.argv or 'reportcsv' in sys.argv)

type = newincident_group.add_argument("-type", dest="type", help="Select type slug",
required='newincident' in sys.argv or 'reportcsv' in sys.argv)

newincident_group.add_argument("-address", dest="address", help="IP or Domain", required='newincident' in sys.argv)

evidence_group = newincident_group.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)

init_group = parser.add_argument_group('init arguments')

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

csv_group = parser.add_argument_group('reportcsv arguments')

csv_group._group_actions.append(feed)
csv_group._group_actions.append(type)

csv_group.add_argument("-path", dest="pathcsv",
                    help="Set Path of CSV file", required='reportcsv' in sys.argv)
csv_group.add_argument("-addressheader", dest="address_header",
                    help="Set CSV Address field", required='reportcsv' in sys.argv)
csv_group.add_argument("-delimiter", dest="delimiter",
                    help="(Optional) Set delimiter for CSV. For TAB use $'\\t'", default=None)


#reportFromPathCSV
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)

def addFromCSV(ngen, csv_path, incident_feed, incident_type, address_header, delimiter):
    ngen.reportFromPathCSV(csv_path, incident_feed, incident_type, address_header, delimiter= delimiter)


def getpyngen():
    try:
        apikey = open(ngen_dir + '/apikey').read()
        url = open(ngen_dir + '/url').read()
    except:
        print("Please use init first")
        exit(2)
    return PyNgen(url, apikey)




if (args.revs == "newincident"):
    ngen = getpyngen()
    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)

elif (args.revs == "reportcsv"):
    ngen = getpyngen()
    addFromCSV(ngen, args.pathcsv, args.feed, args.type, args.address_header, args.delimiter)
