nyfthon : suite

bin/nyfthon

@@ -2,9 +2,11 @@
 from nftables import Nftables
 from os import path
 import sys
+import json
 
 NFT = Nftables()
 NAME = path.basename(__file__)
+TABLE_NAME = "inet filter"
 
 def error(message, code):
     print("%s: %s" % (NAME, message), file = sys.stderr)
@@ -16,6 +18,30 @@ def error_no_arg(arg_name):
 def error_invalid_arg(arg_name, value):
     error('invalid %s "%s"' % (arg_name, value), 2)
 
+def usage():
+    print("""
+{name}: dumb Python wrapper for nftables o_O
+usage:
+    {name} list                                - show currently allowed ports
+    {name} help                                - show this :)
+    {name} allow|deny <protocol> <port number> - add or remove a port from the filtered sets
+""".format(name = NAME).strip())
+
+def get_nft_set_for(protocol):
+    return json.loads(NFT.cmd("list set %s allowed_%s" % (TABLE_NAME, protocol))[1])["nftables"][1]["set"]
+
+def show_ports():
+    NFT.set_json_output(True)
+
+    tcp_ports = get_nft_set_for("tcp")
+    udp_ports = get_nft_set_for("udp")
+    
+    if "elem" in tcp_ports:
+        print("Allowed TCP ports : %s" % (tcp_ports["elem"]))
+
+    if "elem" in udp_ports:
+        print("Allowed UDP ports: %s" % (udp_ports["elem"]))
+
 def allowed_ports(operation, protocol, port):
     if operation != "add" and operation != "delete":
         raise ValueError("not a valid nftables operation")
@@ -26,7 +52,7 @@ def allowed_ports(operation, protocol, port):
     if port < 0 or port > 65535:
         raise ValueError("not a valid port number")
 
-    NFT.cmd("%s element inet filter allowed_%s { %d }" % (operation, protocol, port))
+    NFT.cmd("%s element %s allowed_%s { %d }" % (operation, TABLE_NAME, protocol, port))
 
 def manage_ports(action, args):
     match action:
@@ -70,5 +96,7 @@ if __name__ == "__main__":
             usage()
         case "allow" | "deny":
             manage_ports(action, sys.argv)
+        case "list":
+            show_ports()
         case _:
             error_invalid_arg("action", action)