#!/usr/bin/env python3
from sys import argv
from sys import stderr
from sys import exit
from os import path
import nftables

NFT = nftables.Nftables()
FILTER_TABLE = "inet filter"
try:
    NAME = path.basename(argv.pop(0))
except IndexError:
    pass

def error(message, exit_code):
    print("%s: %s" % (NAME, message), file = stderr)
    exit(exit_code)

def alter_set(operation, protocol, port):
    NFT.cmd("%s element %s allowed_%s { %s }" % (operation, FILTER_TABLE, protocol, port))

def public_allow(argv):
    try:
        protocol = argv.pop(0)
    except IndexError:
        error("no protocol supplied", 1)

    try:
        port = argv.pop(0)
    except IndexError:
        error("no port supplied", 1)

    alter_set('add', protocol, port)

def public_deny(argv):
    try:
        protocol = argv.pop(0)
    except IndexError:
        error("no protocol supplied", 1)

    try:
        port = argv.pop(0)
    except IndexError:
        error("no port supplied", 1)

    alter_set('delete', protocol, port)

if __name__ == '__main__':
    try:
        arg = argv.pop(0)
    except IndexError:
        error("No command supplied", 1)

    try:
        command = globals()["public_%s" % (arg)]
    except KeyError:
        error('invalid command "%s"' % (arg), 2)

    command(argv)