#!/usr/bin/python

#
# color_changer.py
#
# Used to post-process the output of mptcpplot run with the -t flag.
# Changes tokens (Ex: TOKEN001, TOKEN002) to colors in the xplot files
# generated by mptcpplot. Each run of color_changer.py should operate
# only on one set of files (mapping file and xplot files) generated
# by mptcpplot. These files by default will have the same prefix
# when they are output by mptcpplot.
# 
# Output: xplot_file1.new [xplot_file2.new ...]
#         xplot.org compatible files with colors rather than tokens.
#
# Usage:
# color_changer master_file mapping_file xplot_file1 [xplot_file2 ...]
#
# master_file: text file containing mappings between regular expressions
#              and valid colors in xplot.org. One mapping per line.
#              Example:
#              10.155.0.1:[0-9]+:10.155.0.2:[0-9]+ blue
#              10.155.1.1:[0-9]+:10.155.1.2:[0-9]+ red
#
# mapping_file: xxxxx-MAPPING.txt file generated by mptcpplot.
#
# xplot_file1: Mandatory xplot file that contains TOKENs to be converted
#              to colors.
#
# [xplot_file2 ...]: Optional xplot files that contain TOKENs to be
#                    converted to colors.
#


import sys
import re

# Write an error message and exit the program
def error_and_exit(msg, code):
    sys.stderr.write("ERROR: ")
    sys.stderr.write(msg+"\n")
    sys.exit(code)

# Read the master_file. Lines should contain two columns.
# Col 1: Regular expression in string form. Should be a valid re when passed
#        to re.compile.
# Col 2: String of a valid color in xplot.org.
#
# When finished, token_map contains pairs of regular expressions and
# color strings.
def read_master(master_file, token_map):
    try:
        infile = open(master_file, "r")
    except:
        error_and_exit("Could not open file: %s"%(master_file), 1)

    for line in infile:
        # Deal with blank lines.
        if line == "\n":
            continue
        sline = line.rstrip().split(' ')
        if len(sline) != 2:
            error_and_exit("Invalid line format:\nOffending line:%s"%(line), 1)
        try:
            # Store valid regular expressions and their corresponding colors
            token_map.append((re.compile(sline[0]), sline[1]))
        except:
            error_and_exit("Invalid regular expression.\nOffending line: %s"%(line), 1)            
    infile.close()

# Read mapping file output from mptcpplot. Combine with information
# previously read from the master_file. Lines in the mapping_file contain two
# columns.
#
# Col 1: 4-tuple in the form IP_ADDR:PORT:IP_ADDR:PORT
# Col 2: TOKENxxx string
#
# The end result of this function is the ability to convert TOKENs to colors
# by matching 4-tuples in the MAPPING file with regular expressions from the
# master file. color_map contains regular expressions compiled from TOKENs
# that map to colors.
def read_mapping(mapping_file, token_map, color_map):
    try:
        infile = open(mapping_file, "r")
    except:
        error_and_exit("Could not open file: %s"%(mapping_file), 1)

    for line in infile:
        sline = line.rstrip().split(' ')
        if line == "\n":
            continue
        if len(sline) != 2:
            error_and_exit("Invalid line format:\nOffending line: %s"%(line), 1)
        # mptcpplot run with the -t flag will output columns with TOKENxxx
        # in its xplot and mapping files. If colors are output instead, then
        # mptcpplot was run without the -t flag and needs re-run with the -t
        # flag.
        if not sline[1].startswith("TOKEN"):
            error_and_exit("Encountered non-tokens in mapping file.", 1)
            
        found_match = False
        i = 0
        while not found_match and i < len(token_map):
            result = token_map[i][0].match(sline[0])
            if result:
                try:
                    # compile and store TOKENxxx column along with the
                    # proper color.
                    color_map[re.compile(sline[1])] = token_map[i][1]
                except:
                    error_and_exit("Invalid regular expression: %s"%(sline[1]), 1)
                found_match = True  # break loop after first match
            i += 1
            
    infile.close()

# Open xplot_files one at a time. Replace TOKENxxx tokens with the correct
# color contained in color_map.
def process_xplot_files(xplot_files, color_map):
    for x_file in xplot_files:
        try:
            infile = open(x_file, "r")
        except:
            error_and_exit("Could not open file: %s"%(x_file), 1)
        try:
            outfile = open(x_file+".new", "w")
        except:
            error_and_exit("Could not create output file.", 1)

        for line in infile:
            outline = line
            for entry in color_map:
                outline = re.sub(entry, color_map[entry], outline)
            outfile.write(outline)

        infile.close()
        outfile.close()
    

    
usage_string = "color_changer master_file mapping_file xplot_file1 [xplot_file2 ...]"

if len(sys.argv) < 4:
    error_and_exit("Too few arguments.\nUsage: " + usage_string, 1)

master_file = sys.argv[1]   # created by hand
mapping_file = sys.argv[2]  # created by mptcpplot
xplot_files = sys.argv[3:]  # created by mptcpplot

token_map = []
color_map = {}

read_master(master_file, token_map)
read_mapping(mapping_file, token_map, color_map)
process_xplot_files(xplot_files, color_map)








sys.exit(0)
