#! /usr/bin/env python

# This file is part of IVRE.
# Copyright 2011 - 2020 Pierre LALET <pierre@droids-corp.org>
#
# IVRE is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# IVRE is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU General Public License
# along with IVRE. If not, see <http://www.gnu.org/licenses/>.

"""
Access and query the flows database.

See doc/FLOW.md for more information.
"""


from __future__ import print_function
from argparse import ArgumentParser
import datetime
import os
import sys
try:
    reload(sys)
except NameError:
    pass
else:
    sys.setdefaultencoding('utf-8')


from builtins import input
from future.utils import viewitems
try:
    import matplotlib
    import matplotlib.pyplot as plt
except ImportError:
    plt = None


from ivre.db import db
from ivre import utils, config
import ivre.flow


addr_fields = {
    'src': {'type': 'edges', 'field': 'src.addr'},
    'dst': {'type': 'edges', 'field': 'dst.addr'},
    'host': {'type': 'nodes', 'field': 'addr'}
}


def get_addr_argument(field, value):
    addr_field = addr_fields[field]
    # Detect CIDR
    op = '='
    if '/' in value:
        op = '=~'
    return (addr_field['type'], "%s %s %s" % (addr_field['field'], op, value))


def print_fields():
    equals = "=" * 7
    title = "General"
    sys.stdout.write("{0} {1:^10} {0}\n".format(equals, title))
    sys.stdout.writelines(("%-12s: %s\n" % (field, ivre.flow.FIELDS[field])
                           for field in ivre.flow.FIELDS))
    for meta in ivre.flow.META_DESC:
        sys.stdout.write("{0} {1:^10} {0}\n".format(equals, meta))
        sys.stdout.writelines(("meta.%s.%s (list)\n" % (meta, name)
                               for name in ivre.flow.META_DESC[meta]["keys"]))
        sys.stdout.writelines(("meta.%s.%s\n" % (meta, name)
                               for name in ivre.flow.META_DESC[meta].get(
                                   "counters", [])))


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument('--init', '--purgedb', action='store_true',
                        help='Purge or create and initialize the database.')
    parser.add_argument('--ensure-indexes', action='store_true',
                        help='Create missing indexes (will lock the '
                        'database).')
    parser.add_argument('--node-filters', '-n', nargs="+", metavar="FILTER",
                        help='Filter the results with a list of ivre specific '
                        'node textual filters (see WebUI doc in FLOW.md).')
    parser.add_argument('--flow-filters', '-f', nargs="+", metavar="FILTER",
                        help='Filter the results with a list of ivre specific '
                        'flow textual filters (see WebUI doc in FLOW.md).')
    parser.add_argument('--json', '-j', action='store_true',
                        help='Outputs the full json records of results.')
    parser.add_argument('--count', '-c', action='store_true',
                        help='Only return the count of the results.')
    parser.add_argument('--limit', '-l', type=int, default=None,
                        help='Ouput at most LIMIT results.')
    parser.add_argument('--skip', type=int, default=0,
                        help='Skip first SKIP results.')
    parser.add_argument('--orderby', '-o',
                        help='Order of results ("src", "dst" or "flow")')
    parser.add_argument('--separator', '-s', help="Separator string.")
    parser.add_argument('--top', '-t', nargs="+",
                        help='Top flows for a given set of fields, e.g. '
                        '"--top src.addr dport".')
    parser.add_argument('--collect', '-C', nargs="+",
                        help='When using --top, also collect these '
                        'properties.', default=[])
    parser.add_argument('--sum', '-S', nargs="+",
                        help='When using --top, sum on these properties to '
                        'order the result.', default=[])
    parser.add_argument('--least', '-L', action='store_true',
                        help='When using --top, sort records by least')
    parser.add_argument('--mode', '-m',
                        help="Query special mode (flow_map, talk_map...)")
    parser.add_argument('--timeline', '-T', action="store_true",
                        help='Retrieves the timeline of each flow')
    parser.add_argument('--flow-daily', action="store_true",
                        help="Flow count per times of the day. If --precision "
                        "is absent, it will be based on FLOW_TIME_PRECISION "
                        "(%d)" % config.FLOW_TIME_PRECISION)
    parser.add_argument('--plot', action="store_true",
                        help="Plot data when possible (requires matplotlib).")
    parser.add_argument('--fields', nargs='*',
                        help="Without values, gives the list of available "
                        "fields. Otherwise, display these fields for each "
                        "entry.")
    parser.add_argument('--reduce-precision', type=int,
                        metavar="NEW_PRECISION",
                        help="Only with MongoDB backend. "
                        "Reduce precision to NEW_PRECISION for flows "
                        "timeslots. Uses precision, before, after and "
                        "filters.")
    parser.add_argument("--after", "-a", type=str, help="Only with MongoDB "
                        "backend. Get only flows seen after this date. "
                        "Date format: YEAR-MONTH-DAY HOUR:MINUTE. "
                        "Based on timeslots precision. If the given date is "
                        "in the middle of a timeslot, flows start at the next "
                        "timeslot.")
    parser.add_argument("--before", "-b", type=str, help="Only with MongoDB "
                        "backend. Get only flows seen before this date. "
                        "Date format: YEAR-MONTH-DAY HOUR:MINUTE. "
                        "Based on timeslots precision. If the given date is "
                        "in the middle of a timeslot, the whole period is "
                        "kept even if theoretically some flows may have been "
                        "seen after the given date.")
    parser.add_argument('--precision', nargs='?', default=None, const=0,
                        help="Only With MongoDB backend. If PRECISION is "
                        "specified, get only flows with one timeslot of "
                        "the given precision. Otherwise, list "
                        "precisions.", type=int)
    parser.add_argument('--host', type=str, metavar="HOST", help="Filter on "
                        "source OR destination IP. Accepts IP address or "
                        "CIDR.")
    parser.add_argument('--src', type=str, metavar="SRC", help="Filter on "
                        "source IP. Accepts IP address or CIDR.")
    parser.add_argument('--dst', type=str, metavar="DST", help="Filter on "
                        "destination IP. Accepts IP address or CIDR.")
    parser.add_argument('--proto', type=str, metavar="PROTO", help="Filter on "
                        "transport protocol.")
    parser.add_argument('--tcp', action="store_true", help="Alias to "
                        "--proto tcp")
    parser.add_argument('--udp', action="store_true", help="Alias to "
                        "--proto udp")
    parser.add_argument('--port', type=int, metavar="PORT", help="Alias to "
                        "--dport")
    parser.add_argument("--dport", type=int, metavar="DPORT", help="Filter on "
                        "destination port.")
    parser.add_argument("--sport", type=int, metavar="SPORT", help="Filter on "
                        "source port.")
    args = parser.parse_args()

    out = sys.stdout

    if args.plot and plt is None:
        utils.LOGGER.critical("Matplotlib is required for --plot")
        sys.exit(-1)

    if args.init:
        if os.isatty(sys.stdin.fileno()):
            out.write(
                'This will remove any flow result in your database. '
                'Process ? [y/N] ')
            ans = input()
            if ans.lower() != 'y':
                sys.exit(-1)
        db.flow.init()
        sys.exit(0)

    if args.ensure_indexes:
        if os.isatty(sys.stdin.fileno()):
            out.write(
                'This will lock your database. '
                'Process ? [y/N] ')
            ans = input()
            if ans.lower() != 'y':
                sys.exit(-1)
        db.flow.ensure_indexes()
        sys.exit(0)

    if args.fields is not None and not args.fields:
        # Print fields list
        print_fields()
        sys.exit(0)
    elif args.fields is not None:
        # Validate given fields
        for field in args.fields:
            ivre.flow.validate_field(field)

    if args.precision == 0:
        # Get precisions list
        out.writelines('%d\n' % precision
                       for precision in db.flow.list_precisions())
        sys.exit(0)

    filters = {"nodes": args.node_filters or [],
               "edges": args.flow_filters or []}

    args_dict = vars(args)

    for key in addr_fields:
        if args_dict[key] is not None:
            flt_t, flt_v = get_addr_argument(key, args_dict[key])
            filters[flt_t].append(flt_v)

    if args.proto is not None:
        filters['edges'].append("proto = %s" % args.proto)
    for key in ['tcp', 'udp']:
        if args_dict[key]:
            filters['edges'].append("proto = %s" % key)

    for key in ['port', 'dport']:
        if args_dict[key] is not None:
            filters['edges'].append("dport = %d" % args_dict[key])

    if args.sport is not None:
        filters['edges'].append("ANY sports = %d" % args.sport)

    time_args = ['before', 'after']
    time_values = {}
    for arg in time_args:
        time_values[arg] = (
            datetime.datetime.strptime(args_dict[arg], "%Y-%m-%d %H:%M")
            if args_dict[arg] is not None
            else None)

    query = db.flow.from_filters(filters, limit=args.limit, skip=args.skip,
                                 orderby=args.orderby, mode=args.mode,
                                 timeline=args.timeline,
                                 after=time_values['after'],
                                 before=time_values['before'],
                                 precision=args.precision)

    if args.reduce_precision:
        if os.isatty(sys.stdin.fileno()):
            out.write(
                'This will permanently reduce the precision of your '
                'database. Process ? [y/N] ')
            ans = input()
            if ans.lower() != 'y':
                sys.exit(-1)
        new_precision = args.reduce_precision
        db.flow.reduce_precision(new_precision, flt=query,
                                 before=time_values['before'],
                                 after=time_values['after'],
                                 current_precision=args.precision)
        sys.exit(0)

    sep = args.separator or ' | '
    coma = ' ;' if args.separator else ' ; '
    coma2 = ',' if args.separator else ', '
    if args.count:
        count = db.flow.count(query)
        out.write('%(clients)d clients\n%(servers)d servers\n'
                  '%(flows)d flows\n' % count)

    elif args.top:
        top = db.flow.topvalues(query, args.top, collect_fields=args.collect,
                                sum_fields=args.sum, topnbr=args.limit,
                                skip=args.skip, least=args.least)
        for rec in top:
            sys.stdout.write("%s%s%s%s%s\n" % (
                '(' + coma2.join(str(val) for val in rec["fields"]) + ')',
                sep,
                rec["count"],
                sep,
                coma.join(str('(' +
                              coma2.join(str(val) for val in collected) +
                              ')')
                          for collected in rec["collected"])
                if rec["collected"] else ""
            ))

    elif args.flow_daily:
        precision = (args.precision if args.precision is not None
                     else config.FLOW_TIME_PRECISION)
        plot_data = {}
        for rec in db.flow.flow_daily(precision, query,
                                      after=time_values['after'],
                                      before=time_values['before']):
            out.write(
                sep.join([
                    rec["time_in_day"].strftime("%T.%f"),
                    " ; ".join(['(' + x[0] + ', ' + str(x[1]) + ')'
                                for x in rec["flows"]])]))
            out.write("\n")

            if args.plot:
                for flw in rec["flows"]:
                    t = rec["time_in_day"]
                    # pyplot needs datetime objects
                    dt = datetime.datetime(1970, 1, 1,
                                           hour=t.hour,
                                           minute=t.minute,
                                           second=t.second)
                    plot_data.setdefault(flw[0], {})
                    plot_data[flw[0]][dt] = flw[1]
        if args.plot and plot_data:
            t = datetime.datetime(1970, 1, 1, 0, 0, 0)
            t += datetime.timedelta(seconds=config.FLOW_TIME_BASE % precision)
            times = []
            while t < datetime.datetime(1970, 1, 2):
                times.append(t)
                t = t + datetime.timedelta(seconds=precision)
            ax = plt.subplots()[1]
            fmt = matplotlib.dates.DateFormatter('%H:%M:%S')
            for flow, data in viewitems(plot_data):
                values = [(data[ti] if ti in data else 0) for ti in times]
                plt.step(times, values, '.-', where='post', label=flow)
            plt.legend(loc='best')
            ax.xaxis.set_major_formatter(fmt)
            plt.gcf().autofmt_xdate()
            plt.show()

    else:
        fmt = '%%s%s%%s%s%%s' % (sep, sep)
        node_width = len('XXXX:XXXX:XXXX:XXXX:XXXX:XXXX')
        flow_width = len('tcp/XXXXX')
        for res in db.flow.to_iter(query, limit=args.limit, skip=args.skip,
                                   orderby=args.orderby, mode=args.mode,
                                   timeline=args.timeline):
            if args.json:
                out.write('%s\n' % res)
            else:
                elts = {}
                for elt in ["src", "flow", "dst"]:
                    elts[elt] = res[elt]['label']
                    if args.fields:
                        elts[elt] = "%s%s%s" % (
                            elts[elt],
                            coma,
                            coma.join(
                                str(res[elt]['data'].get(field, ""))
                                for field in args.fields
                            )
                        )
                src, flow, dst = elts["src"], elts["flow"], elts["dst"]
                node_width = max(node_width, len(src), len(dst))
                flow_width = max(flow_width, len(flow))
                if not args.separator:
                    fmt = ('%%-%ds%s%%-%ds%s%%-%ds' %
                           (node_width, sep, flow_width, sep, node_width))
                out.write(fmt % (src, flow, dst))
                if args.timeline:
                    out.write(sep)
                    # Print '?' instead of failing if meta.times does not exist
                    try:
                        out.write(coma.join(
                            str(elt) for elt in sorted(
                                res['flow']['data']['meta']['times']
                            )
                        ))
                    except KeyError:
                        out.write("?")
                out.write('\n')
