#!/usr/bin/env python3

import ipfix.ie
import ipfix.reader
import argparse
import csv
from sys import stdin, stdout

def init_ipfix(specfiles = None):
    ipfix.ie.use_iana_default()
    ipfix.ie.use_5103_default()
    
    if specfiles:
        for sf in specfiles:
            ipfix.ie.use_specfile(sf)

def stream_to_csv(instream, ienames):
    cols = ipfix.ie.spec_list(ienames)

    r = ipfix.reader.from_stream(instream)
    w = csv.writer(stdout, dialect='unix')

    w.writerow([e.name for e in cols])
    
    for rec in r.records_as_tuple(cols):
        w.writerow([str(val) for val in rec])
    
    
parser = argparse.ArgumentParser(description="Convert an IPFIX file or stream to CSV")
parser.add_argument('iename', metavar="ie", nargs="+",
                    help="column name of IE")
parser.add_argument('--spec', '-s', metavar="specfile", action="append",
                    help="file to load additional IESpecs from")
parser.add_argument('--file', '-f', metavar="file", nargs="?",
                    help="IPFIX file to read (default stdin)")

args = parser.parse_args()

init_ipfix(args.spec)

if args.file:
    with open (args.file, mode="rb") as f:
        stream_to_csv(f, args.iename)
else:
    stdin = stdin.detach()
    stream_to_csv(stdin, args.iename)