Update toriptables2.py

toriptables2.py

@@ -1,5 +1,5 @@
 #! /usr/bin/env python
-# Written by rupe
+# By Rupe
 """
 Tor Iptables script is an anonymizer
 that sets up iptables and tor to route all services
@@ -11,85 +11,85 @@ from commands import getoutput
 from subprocess import call
 from os.path import isfile
 from os import devnull
+from atexit import register
 from argparse import ArgumentParser
 
-fnull = open(devnull, 'w')
-
 
 class TorIptables(object):
-    def __init__(self):
-        self.tor_config_file = '/etc/tor/torrc'
-        self.torrc = '''
+
+  def __init__(self):
+    self.tor_config_file = '/etc/tor/torrc'
+    self.torrc = '''
 VirtualAddrNetwork 10.0.0.0/10
 AutomapHostsOnResolve 1
 TransPort 9040
 DNSPort 53
 '''
-        self.non_tor_net = ["192.168.0.0/16", "172.16.0.0/12"]
-        self.non_tor = ["127.0.0.0/9", "127.128.0.0/10", "127.0.0.0/8"]
-        self.tor_uid = getoutput("id -ur debian-tor")  # Tor user uid
-        self.trans_port = "9040"  # Tor port
-        self.load_iptables_rules.__init__(self)
+    self.non_tor_net = ["192.168.0.0/16", "172.16.0.0/12"]
+    self.non_tor = ["127.0.0.0/9", "127.128.0.0/10", "127.0.0.0/8"]
+    self.tor_uid = getoutput("id -ur debian-tor")  # Tor user uid
+    self.trans_port = "9040"  # Tor port
+    self.load_iptables_rules.__init__(self)
 
-    def flush_iptables_rules(self):
-        call(["iptables", "-F"])
-        call(["iptables", "-t", "nat", "-F"])
+  def flush_iptables_rules(self):
+    call(["iptables", '-F'])
+    call(["iptables", "-t", "nat", "-F"])
 
-    def load_iptables_rules(self):
-        self.flush_iptables_rules()
-        self.non_tor.extend(self.non_tor_net)
+  def load_iptables_rules(self):
+    self.flush_iptables_rules()
+    self.non_tor.extend(self.non_tor_net)
 
-        call(["iptables", "-t", "nat", "-A", "OUTPUT", "-m", "owner",
-              "--uid-owner", "%s" % self.tor_uid, "-j", "RETURN"])
-        call(["iptables", "-t", "nat", "-A", "OUTPUT", "-p", "udp", "--dport",
-              "53", "-j", "REDIRECT", "--to-ports", "53"])
+    @register
+    def restart_tor():
+      fnull = open(devnull, 'w')
+      call(["service", "tor", "restart"], stderr=fnull)
 
-        for self.net in self.non_tor:
-            call(["iptables", "-t", "nat", "-A", "OUTPUT", "-d",
-                  "%s" % self.net, "-j", "RETURN"])
+    call(["iptables", "-t", "nat", "-A", "OUTPUT", "-m", "owner", "--uid-owner",
+          "%s" % self.tor_uid, "-j", "RETURN"])
+    call(["iptables", "-t", "nat", "-A", "OUTPUT", "-p", "udp", "--dport", "53",
+          "-j", "REDIRECT", "--to-ports", "53"])
 
-        call(["iptables", "-t", "nat", "-A", "OUTPUT", "-p", "tcp", "--syn",
-              "-j", "REDIRECT", "--to-ports", "%s" % self.trans_port])
+    for self.net in self.non_tor:
+      call(["iptables", "-t", "nat", "-A", "OUTPUT", "-d", "%s" % self.net,
+            "-j", "RETURN"])
 
-        call(["iptables", "-A", "OUTPUT", "-m", "state", "--state",
-              "ESTABLISHED,RELATED", "-j", "ACCEPT"])
+    call(["iptables", "-t", "nat", "-A", "OUTPUT", "-p", "tcp", "--syn", "-j",
+          "REDIRECT", "--to-ports", "%s" % self.trans_port])
 
-        for self.net in self.non_tor:
-            call(["iptables", "-A", "OUTPUT", "-d", "%s" % self.net, "-j",
-                  "ACCEPT"])
+    call(["iptables", "-A", "OUTPUT", "-m", "state", "--state",
+          "ESTABLISHED,RELATED", "-j", "ACCEPT"])
 
-        call(["iptables", "-A", "OUTPUT", "-m", "owner", "--uid-owner",
-              "%s" % self.tor_uid, "-j", "ACCEPT"])
-        call(["iptables", "-A", "OUTPUT", "-j", "REJECT"])
+    for self.net in (self.non_tor):
+      call(["iptables", "-A", "OUTPUT", "-d", "%s" % self.net, "-j", "ACCEPT"])
 
-        # Restart Tor
-        call(["service", "tor", "restart"], stderr=fnull)
+    call(["iptables", "-A", "OUTPUT", "-m", "owner", "--uid-owner", "%s" %
+          self.tor_uid, "-j", "ACCEPT"])
+    call(["iptables", "-A", "OUTPUT", "-j", "REJECT"])
 
 
 if __name__ == '__main__':
-    parser = ArgumentParser(
-    description='Tor Iptables script for loading and unloading iptables rules')
-    parser.add_argument('-l', '--load',
-                        action='store_true',
-                        help='This option will load tor iptables rules')
-    parser.add_argument('-f', '--flush',
-                        action='store_true',
-                        help='This option flushes the iptables rules to default')
-    args = parser.parse_args()
+  parser = ArgumentParser(
+      description='Tor Iptables script for loading and unloading iptables rules')
+  parser.add_argument('-l', '--load',
+                      action="store_true",
+                      help='This option will load tor iptables rules')
+  parser.add_argument('-f', '--flush',
+                      action='store_true',
+                      help='This option flushes the iptables rules to default')
+  args = parser.parse_args()
 
-    try:
-        load_tables = TorIptables()
-        if isfile(load_tables.tor_config_file):
-            if not 'VirtualAddrNetwork' in open(
-                load_tables.tor_config_file).read():
-                with open(load_tables.tor_config_file, 'a+') as torrconf:
-                    torrconf.write(load_tables.torrc)
+  try:
+    load_tables = TorIptables()
+    if isfile(load_tables.tor_config_file):
+      if not 'VirtualAddrNetwork' in open(load_tables.tor_config_file).read():
+        with open(load_tables.tor_config_file, 'a+') as torrconf:
+          torrconf.write(load_tables.torrc)
 
-        if args.load:
-            load_tables.load_iptables_rules()
-        elif args.flush:
-            load_tables.flush_iptables_rules()
-        else:
-            parser.print_help()
-    except Exception as err:
-        print(err)
+    if args.load:
+      load_tables.load_iptables_rules()
+    elif args.flush:
+      load_tables.flush_iptables_rules()
+    else:
+      parser.print_help()
+  except Exception as err:
+    print(err)