#!/usr/bin/env python
#
# Naxsi learning utility
#

# Builtins
import glob, fcntl, termios
import sys
import socket
import time
import os
import tempfile
import subprocess
import json
import logging
from collections import defaultdict
from optparse import OptionParser, OptionGroup

# Third party
import elasticsearch
from nxapi.nxtransform import *
from nxapi.nxparse import *

# Default values
F_SETPIPE_SZ = 1031  # Linux 2.6.35+
F_GETPIPE_SZ = 1032  # Linux 2.6.35+

# Initialize logging
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                    format=None,
                    datefmt=None)

def open_fifo(fifo):
    try:
        os.mkfifo(fifo)
    except OSError:
        logging.warning("Fifo ["+fifo+"] already exists (non fatal).")
    except Exception, e:
        logging.error("Unable to create fifo ["+fifo+"]")
    try:
        logging.debug("Opening fifo ... will return when data is available.")
        fifo_fd = open(fifo, 'r')
        fcntl.fcntl(fifo_fd, F_SETPIPE_SZ, 1000000)
        logging.debug("Pipe (modified) size : "+str(fcntl.fcntl(fifo_fd, F_GETPIPE_SZ)))
    except Exception, e:
        logging.error("Unable to create fifo, error: "+str(e))
        return None
    return fifo_fd

def macquire(line):
    z = parser.parse_raw_line(line)
    # add data str and country
    if z is not None:
        for event in z['events']:
            event['date'] = z['date']
            try:
                event['coords'] = geoloc.ip2ll(event['ip'])
                event['country'] = geoloc.ip2cc(event['ip'])
            except NameError:
                pass
        logging.debug(z)
        injector.insert(z)
    else:
        pass

opt = OptionParser()
# group : config
p = OptionGroup(opt, "Configuration options")
p.add_option('-c', '--config', dest="cfg_path", default="/usr/local/etc/nxapi.json", help="Path to nxapi.json (config).")
p.add_option('--colors', dest="colors", action="store_false", default="true", help="Disable output colorz.")
# p.add_option('-q', '--quiet', dest="quiet_flag", action="store_true", help="Be quiet.")
# p.add_option('-v', '--verbose', dest="verb_flag", action="store_true", help="Be verbose.")
opt.add_option_group(p)
# group : in option
p = OptionGroup(opt, "Input options (log acquisition)")
p.add_option('--files', dest="files_in", help="Path to log files to parse.")
p.add_option('--fifo', dest="fifo_in", help="Path to a FIFO to be created & read from. [infinite]")
p.add_option('--stdin', dest="stdin", action="store_true", help="Read from stdin.")
p.add_option('--no-timeout', dest="infinite_flag", action="store_true", help="Disable timeout on read operations (stdin/fifo).")
p.add_option('--syslog', dest="syslog_in", action="store_true", help="Listen on tcp port for syslog logging.")
opt.add_option_group(p)
# group : filtering
p = OptionGroup(opt, "Filtering options (for whitelist generation)")
p.add_option('-s', '--server', dest="server", help="FQDN to which we should restrict operations.")
p.add_option('--filter', dest="filter", action="append", help="This option specify a filter for each type of filter, filter are merge with existing templates/filters. (--filter 'uri /foobar')")
opt.add_option_group(p)
# group : tagging
p = OptionGroup(opt, "Tagging options (tag existing events in database)")
p.add_option('-w', '--whitelist-path', dest="wl_file", help="A path to whitelist file, will find matching events in DB.")
p.add_option('-i', '--ip-path', dest="ips", help="A path to IP list file, will find matching events in DB.")
p.add_option('--tag', dest="tag", action="store_true", help="Actually tag matching items in DB.")
opt.add_option_group(p)
# group : whitelist generation
p = OptionGroup(opt, "Whitelist Generation")
p.add_option('-f', '--full-auto', dest="full_auto", action="store_true", help="Attempt fully automatic whitelist generation process.")
p.add_option('-t', '--template', dest="template", help="Path to template to apply.")
p.add_option('--slack', dest="slack", action="store_false", help="Enables less strict mode.")
p.add_option('--type', dest="type_wl", action="store_true", help="Generate whitelists based on param type")
opt.add_option_group(p)
# group : statistics
p = OptionGroup(opt, "Statistics Generation")
p.add_option('-x', '--stats', dest="stats", action="store_true", help="Generate statistics about current's db content.")
opt.add_option_group(p)
# group : interactive generation
p = OptionGroup(opt, "Interactive Whitelists Generation")
p.add_option('-g', '--interactive-generation', dest="int_gen", action="store_true", help="Use your favorite text editor for whitelist generation.")
opt.add_option_group(p)

(options, args) = opt.parse_args()


try:
    cfg = NxConfig(options.cfg_path)
except ValueError:
    sys.exit(-1)

if options.server is not None:
    cfg.cfg["global_filters"]["server"] = options.server

# https://github.com/nbs-system/naxsi/issues/231
mutally_exclusive = ['stats', 'full_auto', 'template', 'wl_file', 'ips', 'files_in', 'fifo_in', 'syslog_in']
count=0
for x in mutally_exclusive:
    if options.ensure_value(x, None) is not None:
        count += 1
if count > 1:
    logging.critical("Mutually exclusive options are present (ie. import and stats), aborting.")
    sys.exit(-1)


cfg.cfg["output"]["colors"] = "false" if options.int_gen else str(options.colors).lower()
cfg.cfg["naxsi"]["strict"] = str(options.slack).lower()

def get_filter(arg_filter):
    x = {}
    to_parse = []
    kwlist = ['server', 'uri', 'zone', 'var_name', 'ip', 'id', 'content', 'country', 'date',
              '?server', '?uri', '?var_name', '?content']
    try:
        for argstr in arg_filter:
            argstr = ' '.join(argstr.split())
            to_parse += argstr.split(' ')
        if [a for a in kwlist if a in to_parse]:
            for kw in to_parse:
                if kw in kwlist:
                    x[kw] = to_parse[to_parse.index(kw)+1]
        else:
            raise
    except:
        logging.critical('option --filter must have at least one option')
        sys.exit(-1)
    return x

if options.filter is not None:
    cfg.cfg["global_filters"].update(get_filter(options.filter))

try:
    use_ssl = bool(cfg.cfg["elastic"]["use_ssl"])
except KeyError:
    use_ssl = False

es = elasticsearch.Elasticsearch(cfg.cfg["elastic"]["host"], use_ssl=use_ssl)
# Get ES version from the client and avail it at cfg
es_version =  es.info()['version'].get('number', None)
if es_version is not None:
    cfg.cfg["elastic"]["version"] = es_version.split(".")[0]
if cfg.cfg["elastic"].get("version", None) is None:
    logging.critical("Failed to get version from ES, Specify version ['1'/'2'/'5'] in [elasticsearch] section")
    sys.exit(-1)

translate = NxTranslate(es, cfg)



if options.type_wl is True:
    translate.wl_on_type()
    sys.exit(0)

# whitelist generation options
if options.full_auto is True:
    translate.load_cr_file(translate.cfg["naxsi"]["rules_path"])
    results = translate.full_auto()
    if results:
        for result in results:
            logging.info("{0}".format(result))
    else:
        logging.critical("No hits for this filter.")
        sys.exit(1)
    sys.exit(0)

if options.template is not None:
    scoring = NxRating(cfg.cfg, es, translate)

    tpls = translate.expand_tpl_path(options.template)
    gstats = {}
    if len(tpls) <= 0:
        logging.error("No template matching")
        sys.exit(1)
    # prepare statistics for global scope
    scoring.refresh_scope('global', translate.tpl2esq(cfg.cfg["global_filters"]))
    for tpl_f in tpls:
        scoring.refresh_scope('rule', {})
        scoring.refresh_scope('template', {})

        logging.debug(translate.grn.format("Loading template '"+tpl_f+"'"))
        tpl = translate.load_tpl_file(tpl_f)
        # prepare statistics for filter scope
        scoring.refresh_scope('template', translate.tpl2esq(tpl))
        logging.debug("Hits of template : "+str(scoring.get('template', 'total')))

        whitelists = translate.gen_wl(tpl, rule={})
        logging.debug(str(len(whitelists))+" whitelists ...")
        for genrule in whitelists:
            #pprint.pprint(genrule)
            scoring.refresh_scope('rule', genrule['rule'])
            scores = scoring.check_rule_score(tpl)
            if (len(scores['success']) > len(scores['warnings']) and scores['deny'] == False) or cfg.cfg["naxsi"]["strict"] == "false":
                logging.debug(translate.fancy_display(genrule, scores, tpl))
                logging.debug(translate.grn.format(translate.tpl2wl(genrule['rule'], tpl)).encode('utf-8'))
    sys.exit(0)

# tagging options

if options.wl_file is not None and options.server is None:
    logging.critical(translate.red.format("Cannot tag events in database without a server name !"))
    sys.exit(2)

if options.wl_file is not None:
    wl_files = []
    wl_files.extend(glob.glob(options.wl_file))
    count = 0
    for wlf in wl_files:
        logging.debug(translate.grn.format("Loading template '"+wlf+"'"))
        try:
            wlfd = open(wlf, "r")
        except:
            logging.critical(translate.red.format("Unable to open wl file '"+wlf+"'"))
            sys.exit(-1)
        for wl in wlfd:
            [res, esq] = translate.wl2esq(wl)
            if res is True:
                count = 0
                while True:
                    x = translate.tag_events(esq, "Whitelisted", tag=options.tag)
                    count += x
                    if x == 0:
                        break

        logging.debug(translate.grn.format(str(count)) + " items tagged ...")
        count = 0
    sys.exit(0)

if options.ips is not None:
    ip_files = []
    ip_files.extend(glob.glob(options.ips))
    tpl = {}
    count = 0
#    esq = translate.tpl2esq(cfg.cfg["global_filters"])

    for wlf in ip_files:
        try:
            wlfd = open(wlf, "r")
        except:
            logging.critical("Unable to open ip file '"+wlf+"'")
            sys.exit(-1)
        for wl in wlfd:
            logging.debug("=>"+wl)
            tpl["ip"] = wl.strip('\n')
            esq = translate.tpl2esq(tpl)
            pprint.pprint(esq)
            pprint.pprint(tpl)
            count += translate.tag_events(esq, "BadIPS", tag=options.tag)
        logging.debug(translate.grn.format(str(count)) + " items to be tagged ...")
        count = 0
    sys.exit(0)

# statistics
if options.stats is True:
    logging.info(translate.red.format("# Whitelist(ing) ratio :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "whitelisted", limit=2):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    logging.info(translate.red.format("# Top servers :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "server", limit=10):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    logging.info(translate.red.format("# Top URI(s) :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "uri", limit=10):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    logging.info(translate.red.format("# Top Zone(s) :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "zone", limit=10):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    logging.info(translate.red.format("# Top Peer(s) :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "ip", limit=10):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    logging.info(translate.red.format("# Top Country(ies) :"))
    for e in translate.fetch_top(cfg.cfg["global_filters"], "country", limit=10):
        try:
            list_e = e.split()
            logging.info('# {0} {1} {2}{3}'.format(translate.grn.format(list_e[0]), list_e[1], list_e[2], list_e[3]))
        except:
            logging.warning("--malformed--")
    sys.exit(0)


def write_generated_wl(filename, results):

    with open('/tmp/{0}'.format(filename), 'w') as wl_file:
        for result in results:
            for key, items in result.iteritems():
                if items:
                    logging.debug("{} {}".format(key, items))
                    if key == 'genrule':
                        wl_file.write("# {}\n{}\n".format(key, items))
                    else:
                        wl_file.write("# {} {}\n".format(key, items))
        wl_file.flush()

def ask_user_for_server_selection(editor, welcome_sentences, selection):
    with tempfile.NamedTemporaryFile(suffix='.tmp') as temporary_file:
        top_selection = translate.fetch_top(cfg.cfg["global_filters"],
                            selection,
                            limit=10
                        )
        temporary_file.write(welcome_sentences)
        for line in top_selection:
            temporary_file.write('{0}\n'.format(line))
        temporary_file.flush()
        subprocess.call([editor, temporary_file.name])
        temporary_file.seek(len(welcome_sentences))
        ret = []
        for line in temporary_file:
            if not line.startswith('#'):
                ret.append(line.strip().split()[0])
    return ret

def ask_user_for_selection(editor, welcome_sentences, selection, servers):
    regex_message = "# as in the --filter option you can add ? for regex\n"
    ret = {}
    for server in servers:
        server_reminder = "server: {0}\n\n".format(server)
        ret[server] = []
        with tempfile.NamedTemporaryFile(suffix='.tmp') as temporary_file:
            temporary_file.write(welcome_sentences + regex_message + server_reminder)
            cfg.cfg["global_filters"]["server"] = server
            top_selection = translate.fetch_top(cfg.cfg["global_filters"],
                                selection,
                                limit=10
                            )
            for line in top_selection:
                temporary_file.write('{0} {1}\n'.format(selection, line))
            temporary_file.flush()
            subprocess.call([editor, temporary_file.name])
            temporary_file.seek(len(welcome_sentences) + len(server_reminder) + len(regex_message))
            for line in temporary_file:
                if not line.startswith('#'):
                    res = line.strip().split()
                    ret[server].append((res[0], res[1]))
    return ret

def generate_wl(selection_dict):
    for key, items in selection_dict.iteritems():
        if not items:
            return False
        global_filters_context = cfg.cfg["global_filters"]
        global_filters_context["server"] = key
        for idx, (selection, item) in enumerate(items):
           global_filters_context[selection] = item
           translate.cfg["global_filters"] = global_filters_context
           logging.debug('generating wl with filters {0}'.format(global_filters_context))
           wl_dict_list = []
           res = translate.full_auto(wl_dict_list)
           del global_filters_context[selection]
           write_generated_wl(
               "server_{0}_{1}.wl".format(
                                    key,
                                    idx if (selection == "uri") else "zone_{0}".format(item),
                                ),
               wl_dict_list
           )

if options.int_gen is True:
    editor = os.environ.get('EDITOR', 'vi')

    welcome_sentences = '{0}\n{1}\n'.format(
        '# all deleted line or starting with a # will be ignore',
        '# if you want to use slack option you have to specify it on the command line options'
    )

    servers = ask_user_for_server_selection(editor, welcome_sentences, "server")

    uris = ask_user_for_selection(editor, welcome_sentences, "uri", servers)
    zones = ask_user_for_selection(editor, welcome_sentences, "zone", servers)

    if uris:
        generate_wl(uris)
    if zones:
        generate_wl(zones)
    # in case the user let uri and zone files empty generate wl for all
    # selected server(s)
    if not uris and not zones:
        for server in servers:
            translate.cfg["global_filters"]["server"] = server
            logging.debug('generating with filters: {0}'.format(translate.cfg["global_filters"]))
            res = translate.full_auto()
            writing_generated_wl("server_{0}.wl".format(server), res)

    sys.exit(0)

# input options, only setup injector if one input option is present
if options.files_in is not None or options.fifo_in is not None or options.stdin is not None or options.syslog_in is not None:
    if options.fifo_in is not None or options.syslog_in is not None:
        injector = ESInject(es, cfg.cfg, auto_commit_limit=1)
    else:
        injector = ESInject(es, cfg.cfg)
    parser = NxParser()
    offset = time.timezone if (time.localtime().tm_isdst == 0) else time.altzone
    offset = offset / 60 / 60 * -1
    if offset < 0:
        offset = str(-offset)
    else:
        offset = str(offset)
    offset = offset.zfill(2)
    parser.out_date_format = "%Y-%m-%dT%H:%M:%S+"+offset #ES-friendly
    try:
        geoloc = NxGeoLoc(cfg.cfg)
    except:
        logging.critical("Unable to get GeoIP")

if options.files_in is not None:
    reader = NxReader(macquire, lglob=[options.files_in])
    reader.read_files()
    injector.stop()
    sys.exit(0)

if options.fifo_in is not None:
    fd = open_fifo(options.fifo_in)
    if options.infinite_flag is True:
        reader = NxReader(macquire, fd=fd, stdin_timeout=None)
    else:
        reader = NxReader(macquire, fd=fd)
    while True:
        logging.debug("start-",)
        if reader.read_files() == False:
            break
        logging.debug("stop")
    logging.debug('End of fifo input...')
    injector.stop()
    sys.exit(0)

if options.syslog_in is not None:
    sysloghost = cfg.cfg["syslogd"]["host"]
    syslogport = cfg.cfg["syslogd"]["port"]
    while 1:
      reader = NxReader(macquire, syslog=True, syslogport=syslogport, sysloghost=sysloghost)
      reader.read_files()
    injector.stop()
    sys.exit(0)

if options.stdin is True:
    if options.infinite_flag:
        reader = NxReader(macquire, lglob=[], fd=sys.stdin, stdin_timeout=None)
    else:
        reader = NxReader(macquire, lglob=[], fd=sys.stdin)
    while True:
        logging.debug("start-",)
        if reader.read_files() == False:
            break
        logging.debug("stop")
    logging.debug('End of stdin input...')
    injector.stop()
    sys.exit(0)

opt.print_help()
sys.exit(0)
