#!python
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 newevent -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=[
                    "newevent", "init", "reportcsv"], help='Set action: %(choices)s')

newevent_group = parser.add_argument_group('newevent arguments')

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

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

newevent_group.add_argument("-notes", dest="notes", help="Notes for event", 
                            required=False, default=None)

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

newevent_group.add_argument("-evidenceFile", dest="evidenceFile",
                            help="Select path of evidence", default=None)

newevent_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 createNewEvent(pyngen, address, event_feed, event_taxonomy, **kargs):
    try:
        res = pyngen.newEvent(address, event_feed, event_taxonomy, **kargs)
        print("Added event: {}".format(res))
    except NewEventFieldError as e:
        print("Error creating event: \n{}".format(e.detail))
        exit(2)

def addFromCSV(ngen, csv_path, event_feed, event_taxonomy, address_header, delimiter):
    ngen.reportFromPathCSV(csv_path, event_feed, event_taxonomy, 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 in "newevent"):
    ngen = getpyngen()
    createNewEvent(ngen, args.address, args.feed,
                      args.type, evidence_text=args.evidenceText, evidence_file=args.evidenceFile, notes=args.notes)
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)
