Args and Policies

functions/iptables_generator.py

@@ -2,7 +2,9 @@
 
 from typing import Generator
 
-def generate_iptable_rules(addresses: list[dict]) -> Generator[str, dict, None]:
+import argparse
+
+def generate_iptable_rules(addresses: list[dict], args: argparse.Namespace) -> Generator[str, dict, None]:
     # Commands
     sudo: str = "sudo"
     iptables: str = "iptables"
@@ -10,7 +12,7 @@ def generate_iptable_rules(addresses: list[dict]) -> Generator[str, dict, None]:
 
     # Variables
     chain_name: str = "PROTECT_FEDI"
-    policy: str = "DROP"  # REJECT tells the server you're dropping them, DROP is more evil in that you drop the connection silently
+    policy: str = args.policy  # REJECT tells the server you're dropping them, DROP is more evil in that you drop the connection silently
 
     # IP Tables Setup
     create_chain: str = f"{sudo} {iptables} -N {chain_name}"

functions/json_formatter.py

@@ -1,7 +1,8 @@
 from typing import Generator
 
 import json
+import argparse
 
-def format_addresses(addresses: list[dict]) -> Generator[str, dict, None]:
+def format_addresses(addresses: list[dict], args: argparse.Namespace) -> Generator[str, dict, None]:
     for address in addresses:
         yield json.dumps(obj=address)

functions/plaintext_formatter.py

@@ -1,6 +1,8 @@
 from typing import Generator
 
-def format_addresses(addresses: list[dict]) -> Generator[str, dict, None]:
+import argparse
+
+def format_addresses(addresses: list[dict], args: argparse.Namespace) -> Generator[str, dict, None]:
     for address in addresses:
         if "route" in address:
             yield address["route"]

main.py

@@ -3,7 +3,7 @@ from functions import plaintext_formatter, whois_lookup, iptables_generator, jso
 import argparse
 
 if __name__ == "__main__":
-    argParser = argparse.ArgumentParser()
+    argParser: argparse.ArgumentParser = argparse.ArgumentParser()
     argParser.add_argument("-f", "--format",
                            default="iptables",
                            const="iptables",
@@ -12,6 +12,14 @@ if __name__ == "__main__":
                            choices=("iptables", "plain", "jsonl"),
                            help="Output format of IP address list (default: %(default)s)")
 
+    argParser.add_argument("-p", "--policy",
+                           default="DROP",
+                           const="DROP",
+                           nargs="?",
+                           type=str,
+                           choices=("DROP", "REJECT", "ACCEPT"),
+                           help="IP tables policy for handling incoming packets (default: %(default)s)")
+
     args = argParser.parse_args()
 
     addresses: list[dict] = []
@@ -23,15 +31,15 @@ if __name__ == "__main__":
     # Generate IP Table Rules
     if args.format == "iptables":
         # IP Tables Commands
-        for rule in iptables_generator.generate_iptable_rules(addresses=addresses):
+        for rule in iptables_generator.generate_iptable_rules(addresses=addresses, args=args):
             print(rule)
     elif args.format == "plain":
         # Just Plain Addresses
-        for address in plaintext_formatter.format_addresses(addresses=addresses):
+        for address in plaintext_formatter.format_addresses(addresses=addresses, args=args):
             print(address)
     elif args.format == "jsonl":
         # JSON Formatted Addresses
-        for address in json_formatter.format_addresses(addresses=addresses):
+        for address in json_formatter.format_addresses(addresses=addresses, args=args):
             print(address)
     else:
         print(f"Unknown format: `{args.format}`")