traffic reports out of ulog-acctd logs

traffic-report.py

#!/usr/bin/env python
'''
/etc/shorewall/rules::

    SECTION ESTABLISHED
    LOG:ULOG loc net
    LOG:ULOG net loc

    SECTION RELATED
    LOG:ULOG loc net
    LOG:ULOG net loc

    SECTION NEW
    LOG:ULOG loc net
    LOG:ULOG net loc


/etc/ulog-acctd.conf::

    multicast groups=1
    accounting file=/var/log/ulog-acctd/account.log
    dump file=/var/log/ulog-acctd/dump
    debug file=/var/log/ulog-acctd/debug.log
    debug = syscall, misc, statistics, error, asdf, error-packet
    accounting format="%t\t%s\t%d\t%b\n"
    empty interface="-"
    empty prefix="-"
    flush=300
    fdelay=0


# apt-get install python-iplib python-prettytable python-argparse

'''
from dateutil.parser import parse as _parse_date
from dateutil.relativedelta import relativedelta as rdelta, MO
from dateutil.tz import tzlocal
from itertools import chain
from prettytable import PrettyTable, FRAME, NONE
import argparse
import csv
import datetime as dt
import gzip
import iplib
import socket
import sys


PERIOD_DAY, PERIOD_WEEK, PERIOD_MONTH = 'day', 'week', 'month'
CHOICES_PERIOD = frozenset([PERIOD_DAY, PERIOD_WEEK, PERIOD_MONTH])

FORMAT_TABLE, FORMAT_CSV = 'table', 'csv'
CHOICES_FORMAT = frozenset([FORMAT_TABLE, FORMAT_CSV])

DEFAULT_FORMAT = FORMAT_TABLE
DEFAULT_PERIOD = PERIOD_DAY
DEFAULT_CIDR_LIST = [iplib.CIDR('192.168.0.0/16')]

FIELD_HOST, FIELD_IN, FIELD_OUT = 'host', '      in', '      out'

UNITS = [
    (1024 ** 4, '%.1fT'),
    (1024 ** 3, '%.1fG'),
    (1024 ** 2, '%.1fM'),
    (1024 ** 1, '%.1fK'),
]

def units(bytes):
    n = float(bytes)
    if n:
        for a, f in UNITS:
            if n >= a:
                return f % (n / a)
    return '%s   ' % bytes


class Table(PrettyTable):

    def get_string(self, start=0, end=None, fields=None,
                   header=True, border=True, hrules=FRAME, sortby=None, reversesort=False):
        string = PrettyTable.get_string(self, start, end, fields, header, border, hrules, sortby, reversesort)
        if header and not start and not end:
            string += '\n' + self._stringify_footer(fields, border, hrules)
            if hrules != NONE:
                string += '\n' + self._stringify_hrule(fields, border)
        return string

    def set_footer(self, footer):
        self.footer = footer

    def _stringify_footer(self, fields=None, border=True, hrules=FRAME):
        return self._stringify_row(self.footer, fields, border, hrules)


class FileType(argparse.FileType):
    def __call__(self, string):
        if string.endswith('.gz'):
            return gzip.open(string)
        return super(FileType, self).__call__(string)


def parse_date(s):
    return _parse_date(s, tzinfos={None: tzlocal()}).date()


def parse_cidr(s):
    cidr = '/' in s and s or '%s/32' % s
    return iplib.CIDR(cidr)


def open_file(f):
    if f.endswith('.gz'):
        return gzip.open(f, 'r')
    return open(f, 'r')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cidr', action='append', dest='cidr_list', type=parse_cidr)
    parser.add_argument('-f', '--format', action='store', default=FORMAT_TABLE, choices=CHOICES_FORMAT)
    parser.add_argument('-p', '--period', default=DEFAULT_PERIOD, choices=CHOICES_PERIOD)
    parser.add_argument('-r', '--resolve', action='store_true', default=False)
    parser.add_argument('-w', '--when', default=dt.date.today(), type=parse_date, metavar='ISO_DATE')
    parser.add_argument('files', type=FileType(), metavar='FILENAME', nargs='*', default=[sys.stdin])
    args = parser.parse_args(sys.argv[1:])

    if args.period == PERIOD_DAY:
        start = args.when
        end = start + rdelta(days=1)
    elif args.period == PERIOD_WEEK:
        start = args.when + rdelta(weekday=MO(-1))
        end = start + rdelta(weeks=1)
    elif args.period == PERIOD_MONTH:
        start = args.when + rdelta(day=1)
        end = start + rdelta(months=1)
    else:
        raise Exception('period argument error')

    cidr_list = args.cidr_list or DEFAULT_CIDR_LIST
    start, end = int(start.strftime('%s')), int(end.strftime('%s'))

    data = {}
    for s in chain(*list([f.readlines() for f in args.files])):
        ts, src, dst, bytes = s.split()
        ts, bytes = int(ts), int(bytes)
        if ts < start and (start - ts) > 3600*24 or ts > end:
            break
        if ts >= start and ts <= end:
            ip, din, dout = None, None, None
            if any(iplib.IPv4Address(src) in cidr for cidr in cidr_list):
                ip, din, dout = src, 0, bytes
            elif any(iplib.IPv4Address(dst) in cidr for cidr in cidr_list):
                ip, din, dout = dst, bytes, 0
            if ip:
                data.setdefault(ip, [0, 0])
                data[ip][0] += din
                data[ip][1] += dout

    if args.format == FORMAT_TABLE:
        table = Table([FIELD_HOST, FIELD_IN, FIELD_OUT])
        table.set_field_align(FIELD_HOST, 'l')
        table.set_field_align(FIELD_IN, 'r')
        table.set_field_align(FIELD_OUT, 'r')

        total = [0, 0]
        for ip, row in data.items():
            try:
                host = args.resolve and socket.gethostbyaddr(ip)[0].lower() or ip
            except socket.herror:
                host = ip
            table.add_row([host, units(row[0]), units(row[1])])
            total[0] += row[0]
            total[1] += row[1]

        table.set_footer(['', units(total[0]), units(total[1])])
        table.printt(sortby=FIELD_HOST)
    elif args.format == FORMAT_CSV:
        writer = csv.writer(sys.stdout)
        writer.writerow(['ip', 'in', 'out'])
        for ip, row in data.items():
            row.insert(0, ip)
            writer.writerow(row)


if __name__ == '__main__':
    main()