#!/usr/bin/env python3

import argparse
import matplotlib.pyplot    as plt
import tempfile
import subprocess
import threading
import os
from   os.path              import dirname, basename, isfile, splitext, normpath, expanduser
import sys
import glob
import pickle

# Configuration: Adapt to your system.
T2_ROOT   = '{}/code/tranalyzer2-0.9.2'.format(expanduser('~'))
DATA_PATH = '/mnt/{}/BPB/data_short'.format(os.environ['USER'])

VERBOSE = False

# Automatically derived
T2_PATH    = '{}/tranalyzer2/src/tranalyzer'.format(T2_ROOT)
T2CONF     = '{}/scripts/t2conf/t2conf'.format(T2_ROOT)
T2BUILD    = '{}/autogen.sh'.format(T2_ROOT)
BPBS_PATH  = 'bpbs.data'

REQUIRED_PLUGINS = ['nFrstPkts', 'sslDecode']

def setup_t2():
    options = {
        'NFRST_IAT':     '0',
        'NFRST_MINIATS': '0',
        'NFRST_MINIATU': '50000',
        'NFRST_PLAVE':   '0',
        'NFRST_PKTCNT':  '200'
    }

    # Configure
    for opt, val in options.items():
        subprocess.call([T2CONF, 'nFrstPkts', '-D', f'{opt}={val}'])

    # Build
    subprocess.call([T2BUILD])
    for plugin in REQUIRED_PLUGINS:
        subprocess.call([T2BUILD, plugin])

def definitely_need_setup_t2():
    active_plugins = subprocess.check_output([T2BUILD, '-l'], text=True).splitlines()
    return not all(plugin in active_plugins for plugin in REQUIRED_PLUGINS)

def run_t2(infile, outdir):
    subprocess.call([
        T2_PATH,
        '-r', infile,
        '-w', outdir
    ], stdout=subprocess.DEVNULL)
    return '{}/{}_flows.txt'.format(outdir, splitext(basename(infile))[0])

def get_bpb(flowfile):
    EXTRACT_SCRIPT = '''
        source "$1/scripts/t2utils.sh"
        flowfile="$2"

        cat $flowfile | $TAWK -t -I $flowfile \
            '($dir == "B" && $sslServerName ~ /.googlevideo.com/) {
                print $numBytesSnt, $L2L3L4Pl_Iat_nP
                }' |
            sort -nr        | # sort by byte count, descending
            head -n 1       | # get biggest flow
            cut -d$'\t' -f2 | # get nFrstPkts output (burst)
            tr ';' '\n'     | # one burst per line
            cut -d'_' -f1   | # extract byte count from each burst
            cat
    '''
    proc = subprocess.Popen(['/bin/bash', '-c', EXTRACT_SCRIPT, '', T2_ROOT, flowfile], stdout=subprocess.PIPE)
    bpbs = [int(bpb) for bpb in proc.stdout] if proc.stdout else []
    return bpbs

def extract_bpb(pcap):
    with tempfile.TemporaryDirectory() as tempdir:
        flowfile = run_t2(pcap, tempdir)
        bpb = get_bpb(flowfile)
    return bpb

try:
    print(extract_bpb(glob.glob(f'{DATA_PATH}/*/Train/*.pcap')[0]))
except:
    print('Unable to find any PCAPs in your training data. Are you sure you set DATA_PATH correctly?')
    pass

def thread_t2_bpb(bpbs, index, pcap):
    bpb = extract_bpb(pcap)
    bpbs[index] = bpb

def get_bpbs_for_class(c):
    pcaps_root = '{}/{}/Train'.format(DATA_PATH.rstrip('/'), c)
    pcaps = glob.glob('{}/*.pcap'.format(pcaps_root))
    bpbs = []

    max_pcaps = 6
    chunk_size = max_pcaps # `chunk_size` must divide `max_pcaps`.
    pcaps_chunks = zip(*[iter(pcaps[:max_pcaps])]*chunk_size)
    for chunk in map(list, pcaps_chunks):
        threads_results = [None] * chunk_size
        threads = [threading.Thread(target=thread_t2_bpb, args=(threads_results, i, pcap))
                   for i, pcap in enumerate(chunk)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        for bpbs_c in threads_results:
            bpbs.append(bpbs_c)
    return bpbs

def load_bpbs(path):
    with open(path, 'rb') as fh:
        return pickle.load(fh)

print(get_bpbs_for_class('Avengers'))

def learn(classes, *, write_to=None):
    bpbs = { c: [] for c in classes }
    for c in classes:
        print('Extracting bursts for class {}'.format(c))
        c_bpb = get_bpbs_for_class(c)
        for b in c_bpb:
            # NOTE: Could be set() to improve performance in production, but we're using
            # list() to preserve the order of the bursts and produce more meaningful plots.
            bpbs[c].append(list(b))
        if write_to:
            with open(write_to, 'wb') as fh:
                pickle.dump(bpbs, fh)
    return bpbs

def similarity(xs, ys):
    return sum(x in ys for x in xs)

def nearest_neighbors(x, bpbs, n=3):
    top, sims = [None] * n, [0] * n
    for c in bpbs:
        best_in_class = 0
        j = -1
        for y in bpbs[c]:
            j += 1
            s = similarity(x, y)
            for i in range(n):
                if s > sims[i] and s > best_in_class:
                    sims[i] = s
                    best_in_class = s
                    top[i] = c
                    break
    return list(zip(top, sims))

def parse_args(argv=sys.argv[1:]):
    parser = argparse.ArgumentParser(description='Classifier for YouTube video streams')
    group = parser.add_argument_group('Classifier options')
    modes = set(['learn', 'test'])
    parser.add_argument('mode', metavar='mode', choices=modes, nargs=1,
            help='Mode to run in. Available: ' + ', '.join(modes))
    group.add_argument('-t', '--test', default='', type=str, help='Path to pcap to classify')
    group.add_argument('-f', '--force', action='store_true', help='Overwrite existing files without prompting')
    group.add_argument('-s', '--setup', action='store_true', help='Perform initial setup of tranalyzer')
    return parser.parse_args(argv)

def plot_bursts(bpbs, title=''):
    width = 10
    cols = 3
    rows = (len(bpbs) + cols - 1 ) // cols
    fig = plt.figure(figsize=(width, 3 * rows))
    ax1 = plt.subplot(rows, cols, 1)
    for i, bpb in enumerate(map(list, bpbs), 1):
        plt.subplot(rows, cols, i, sharey=ax1)

        # Optional: Filter out audio bursts <=500K.
        #bpb = [x for x in bpb if (x > 500*1024)]

        # Optional: Sort for alternative visualization perspective.
        #bpb.sort()

        plt.bar(x=range(len(bpb)), height=bpb)
        plt.title(title + ' ' + str(i), fontsize=12, fontweight=0)

    plt.tight_layout()
    plt.show(block=True)

def main(argv=None):
    args = parse_args(argv)

    if args.setup or definitely_need_setup_t2():
        print('Setting up tranalyzer ... ', end='')
        setup_t2()
        print('done')

    # run t2 on all PCAPs and extract BPB set for each class
    if args.mode[0] == 'learn':
        print('Storing learned BPB models to file: {}'.format(BPBS_PATH))
        if isfile(BPBS_PATH) and not args.force:
            print('WARNING: The model file {} already exists. Overwrite it? [yN] '.format(BPBS_PATH), end='')
            if input().lower() != 'y':
                print('Exiting.')
                sys.exit(1)
        classes = [basename(normpath(p)) for p in glob.glob('{}/**/'.format(DATA_PATH), recursive=False)]
        classes = classes[:10]
        print('Found {} classes, building model now ...'.format(len(classes)))
        bpbs = learn(classes, write_to=BPBS_PATH)
        print('Done building model, ready to test now.')

        # NOTE: Collected BPB features are stored in a file and are used when this program is
        # invoked in "test" mode.

        if VERBOSE:
            for c in classes[:3]:
                print('Bursts for class {}:'.format(c))
                plot_bursts(list(bpbs[c]), c) # Plot bursts of whichever class happens to be the first.
    else:
        if not args.test:
            print("Missing path to test PCAP file. See help page.")
            sys.exit(1)

        pcap_test = args.test
        bpbs = load_bpbs(BPBS_PATH)
        bpb_test = extract_bpb(pcap_test)

        print('Bursts of test sample ({}):'.format(basename(pcap_test)))
        plot_bursts([bpb_test], 'Test sample')

        top = nearest_neighbors(bpb_test, bpbs, 3)
        result = top[0][0]
        print('Classification result:  {}'.format(result))
        print('')
        print('(Top 3: {})'.format(top))
        plot_bursts(bpbs[result], result)

if __name__ == '__main__':
    main(['learn'])

    print('-' * 40)

    test_pcap = f'{DATA_PATH}/Fast_and_Furious_six/Test/Fast_and_Furious_six_Train00_40_30.pcap'
    main(['test', '-t', test_pcap])