#!/bin/python3
#
import argparse
from getpass import getpass
import re
import subprocess

def device_by_vlan(n):
    return 'eth0' if n == 10 else 'eth0.%d' % n

def iprange_by_vlan(n):
    return '192.168.0.0/24' if n == 10 else '192.168.%d.0/24' % n

def get_iptables_info():
    global password
    my_cmd = "sudo -S iptables -L -n -v --line-numbers".split()
    cmd = subprocess.run(my_cmd, stdout=subprocess.PIPE, input=password, encoding="ascii",)
    return cmd.stdout

def get_rule_num(chain, comment):
    required_chain = False
    for line in get_iptables_info().split('\n'):
        if line.strip().lower().startswith('chain'):
            required_chain = line.strip().lower().startswith('chain ' + chain.strip().lower())
        if required_chain:
            rule_num = None
            try:
                rule_num = int(line.split()[0])
            except:
                pass
            else:
                s = re.search(r'\/\*(.*?)\*\/', line)
                if s is not None:
                    if s.group(1).strip() == comment:
                        return rule_num



if __name__ == "__main__":
    global password

    prog = 'vlan manager'
    description = 'Add and removes iptable rules for a given vlan'
    epilog = 'Text at the bottom of help'
    parser = argparse.ArgumentParser(prog=prog, description=description, epilog=epilog)
    parser.add_argument('-v', '--vlan_number', type=int, required=True, help="Define the vlan to be used")
    parser.add_argument('-o', '--open_to_internet', action='store_true', help="Toggle rule: Open vlan to the internet")
    parser.add_argument('-l', '--log_to_internet', action='store_true', help="Toggle rule: Log vlan to the internet")

    args = parser.parse_args()

    UID_O = 'VLANM_RULE__OPEN_TO_INTERNET_%d' % args.vlan_number
    UID_L = 'VLANM_RULE__LOG_TO_INTERNET_%d' % args.vlan_number

    if args.open_to_internet:
        password = getpass("password: ")
        print()
        rn = get_rule_num('forward', UID_O)
        if rn is None:
            print("Add rule:", UID_O)
        else:
            print("Remove rule:", UID_O)
        if getpass('Proceed? [Y/n]') in ['', 'y', 'Y']:
            if rn is None:
                my_cmd = ['sudo',  '-S', 'iptables', '-A', 'FORWARD', '-j', 'ACCEPT', '-i', device_by_vlan(args.vlan_number), '-o', device_by_vlan(10), '!', '-d', '192.168.0.0/16', '-m', 'comment', '--comment', UID_O]
                cmd = subprocess.run(my_cmd, stdout=subprocess.PIPE, input=password, encoding="ascii",)
                print(cmd.stdout)
            else:
                my_cmd = ['sudo', '-S', 'iptables', '-D', 'FORWARD', str(rn)]
                cmd = subprocess.run(my_cmd, stdout=subprocess.PIPE, input=password, encoding="ascii",)
                print(cmd.stdout)
    elif args.log_to_internet:
        password = getpass("password: ")
        print()
        rn = get_rule_num('forward', UID_L)
        if rn is None:
            print("Add rule:", UID_L)
        else:
            print("Remove rule:", UID_L)
        if getpass('Proceed? [Y/n]') in ['', 'y', 'Y']:
            if rn is None:
                my_cmd = ['sudo', '-S', 'iptables', '-A', 'FORWARD', '-j', 'LOG', '-i', device_by_vlan(args.vlan_number), '-o', device_by_vlan(10), '!', '-d', '192.168.0.0/16', '-m', 'comment', '--comment', UID_L]
                cmd = subprocess.run(my_cmd, stdout=subprocess.PIPE, input=password, encoding="ascii",)
                print(cmd.stdout)
                print("Use 'sudo ~/bin/ftail /var/log/kern.log' to view the logs.")
            else:
                my_cmd = ['sudo', '-S', 'iptables', '-D', 'FORWARD', str(rn)]
                cmd = subprocess.run(my_cmd, stdout=subprocess.PIPE, input=password, encoding="ascii",)
                print(cmd.stdout)


    else:
        print("************************")
        print("* No action specified! *")
        print("************************")
        print()
        parser.print_help()