package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
)

func check_err(e error) {
	if e != nil {
		panic(e)
	}
}

var PARSE_CHAIN_MATCH_ACTION_PATTERN = regexp.MustCompile("^-A (?P<chain>[^ ]*) (?P<match>.*) -j (?P<action>.*)$")
var PARSE_CHAIN_ACTION_PATTERN = regexp.MustCompile("^-A (?P<chain>[^ ]*) -j (?P<action>.*)$")
var PARSE_SKIP_PATTERN = regexp.MustCompile("^[#*]|^COMMIT")
var PARSE_TABLES = regexp.MustCompile("^:(?P<table>[^ ]*)")

var TABLES []string

const HEADER = "\033[95m"
const BLUE = "\033[94m"
const GREEN = "\033[92m"
const WARNING = "\033[93m"
const FAIL = "\033[91m"
const ENDC = "\033[0m"
const BOLD = "\033[1m"
const UNDERLINE = "\033[4m"

func customTable(action string) bool {
	for _, name := range TABLES {
		if name == action {
			return true
		}
	}
	return false
}

func colorAction(action string) string {
	if !customTable(action) {
		return fmt.Sprintf("%s%s%s", GREEN, action, ENDC)
	} else {
		return fmt.Sprintf("%s%s%s", BLUE, action, ENDC)
	}
}

func matchPrint(input []string) {
	for _, line := range input {
		if PARSE_SKIP_PATTERN.MatchString(line) {
			//fmt.Println("Skip", line)
			continue
		}

		match := PARSE_TABLES.FindStringSubmatch(line)
		if match != nil {
			//fmt.Println("Table", match[1])
			TABLES = append(TABLES, match[1])
			continue
		}

		match = PARSE_CHAIN_MATCH_ACTION_PATTERN.FindStringSubmatch(line)
		if match != nil {
			//fmt.Println("Chain Match Action", line)
			chain := match[1]
			rule := match[2]
			action := match[3]
			action = colorAction(action)
			fmt.Printf("%s -> %s \n\tif %s\n", chain, action, rule)
			continue
		}

		match = PARSE_CHAIN_ACTION_PATTERN.FindStringSubmatch(line)
		if match != nil {
			//fmt.Println("Chain Action", line)
			chain := match[1]
			action := match[2]
			action = colorAction(action)
			fmt.Printf("%s -> %s\n", chain, action)
			continue
		}

		fmt.Println("Not Matched", line)
	}
}

func main() {
	f, e := os.Open("/tmp/dump.txt")
	check_err(e)

	defer f.Close()
	scanner := bufio.NewScanner(f)
	scanner.Split(bufio.ScanLines)

	var input []string

	for scanner.Scan() {
		input = append(input, scanner.Text())
	}

	matchPrint(input)
}

import re

# sudo iptables-save [-t nat]> /tmp/dump.txt
input = open('/tmp/dump.txt')
PARSE_CHAIN_MATCH_ACTION_PATTERN = re.compile('^-A ([^ ]*) (.*) -j (.*)$')
PARSE_CHAIN_ACTION_PATTERN = re.compile('^-A ([^ ]*) -j (.*)$')
PARSE_SKIP_PATTERN = re.compile('^[#*]|^COMMIT')
PARSE_TABLES = re.compile('^:([^ ]*)')

TABLES = []

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    YELLOW_LINE = WARNING+'{}'+ENDC
    BLUE_LINE = OKBLUE+'{}'+ENDC
    GREEN_LINE = OKGREEN+'{}'+ENDC
    RED_LINE = FAIL+'{}'+ENDC

C = bcolors

def colorize(jump):
    if jump not in TABLES:
        return C.BLUE_LINE.format(jump)
    else:
        return C.GREEN_LINE.format(jump)

def match_print(input):
    for line in input:
        match = re.match(PARSE_SKIP_PATTERN, line)
        if match:
            continue

        match = re.match(PARSE_TABLES, line)
        if match:
            table_name = match.groups()[0]
            TABLES.append(table_name)
            continue

        match = re.match(PARSE_CHAIN_MATCH_ACTION_PATTERN, line)
        if match:
            chain = match.groups()[0]
            jump = match.groups()[-1]
            rule_match = match.groups()[1]
            jump = colorize(jump)
            print("{} -> {} \n\tif {}".format(chain, jump, rule_match))
            continue

        match = re.match(PARSE_CHAIN_ACTION_PATTERN, line)
        if match:
            chain = match.groups()[0]
            jump = match.groups()[-1]
            jump = colorize(jump)
            print("{} -> {}".format(chain, jump))
            continue

        print('NOT PARSED', line)

if __name__ == '__main__':
    match_print(input)