Update yunohost_firewall.py

Refactoring

yunohost_firewall.py

@@ -11,8 +11,8 @@ except ImportError:
 
 
 
-def firewall_allow(protocol=None,port=None,ip=None):
+def firewall_allow(protocol=None,port=None,ipv6=None):
-    if ip == True:
+    if ipv6 == True:
         ip = 'ipv6'
         iptables="ip6tables"
     else:
@@ -21,23 +21,24 @@ def firewall_allow(protocol=None,port=None,ip=None):
 
     if protocol == "Both":
         TCP_rule = iptables+" -A INPUT -p tcp -i eth0 --dport "+ port +" -j ACCEPT"       
-        
         UDP_rule = iptables+" -A INPUT -p udp -i eth0 --dport "+ port +" -j ACCEPT"
-        append_remove_port(port,'tcp','a',ip)
+        
-        append_remove_port(port,'udp','a',ip)
+        update_yml(port,'tcp','a',ip)
+        update_yml(port,'udp','a',ip)
+        
         os.system(TCP_rule)
         os.system(UDP_rule)
 
     else:
         rule = iptables+" -A INPUT -p "+ protocol +" -i eth0 --dport "+ port +" -j ACCEPT"
-        append_remove_port(port,protocol,'a',ip)
+        update_yml(port,protocol,'a',ip)
         os.system(rule)
 
 
 
-def firewall_disallow(protocol=None,port=None,ip=None):
+def firewall_disallow(protocol=None,port=None,ipv6=None):
 
-    if ip == True:
+    if ipv6 == True:
         ip = 'ipv6'
         iptables="ip6tables"
     else:
@@ -47,13 +48,16 @@ def firewall_disallow(protocol=None,port=None,ip=None):
     if protocol == "Both":
         TCP_rule = iptables+" -A INPUT -p tcp -i eth0 --dport "+ port +" -j REJECT"
         UDP_rule = iptables+" -A INPUT -p udp -i eth0 --dport "+ port +" -j REJECT"
-        append_remove_port(port,'tcp','r',ip)
+        
-        append_remove_port(port,'udp','r',ip)
+        update_yml(port,'tcp','r',ip)
+        update_yml(port,'udp','r',ip)
+        
         os.system(TCP_rule)
         os.system(UDP_rule)
+        
     else:
         rule = iptables+" -A INPUT -p "+ protocol +" -i eth0 --dport "+ port +" -j REJECT"
-        append_remove_port(port,protocol,'r',ip)
+        update_yml(port,protocol,'r',ip)
         os.system(rule)
 
 
@@ -64,27 +68,7 @@ def firewall_list():
     '''
     with open ('firewall.yml') as f:
         firewall = yaml.load(f)
-        TCP_port_list_ipv4 = firewall['ipv4']['TCP']
+    return firewall
-        UDP_port_list_ipv4 = firewall['ipv4']['UDP']
-        TCP_port_list_ipv6 = firewall['ipv6']['TCP']
-        UDP_port_list_ipv6 = firewall['ipv6']['UDP']
-        print("Port TCP Open for ipv4:")
-
-        for i,port in enumerate (TCP_port_list_ipv4):
-            print("-"+str(port))
-
-        print("Port UDP Open for ipv4 :")
-        for i,port in enumerate (UDP_port_list_ipv4):
-            print("-"+str(port))
-
-        print("Port TCP Open for ipv6:")
-        for i,port in enumerate (TCP_port_list_ipv6):
-            print("-"+str(port))
-
-        print("Port UDP Open for ipv6 :")
-        for i,port in enumerate (UDP_port_list_ipv6):
-            print("-"+str(port))
-    f.close()
 
 
 
@@ -99,39 +83,32 @@ def firewall_reload():
     with open('firewall.yml','r') as f:
         firewall = yaml.load(f)
 
-        TCP_port_list_ipv4 = firewall['ipv4']['TCP']
-        UDP_port_list_ipv4 = firewall['ipv4']['UDP']
-
-        TCP_port_list_ipv6 = firewall['ipv6']['TCP']
-        UDP_port_list_ipv6 = firewall['ipv6']['UDP']
-    f.close()
-
     os.system ("iptables -P INPUT ACCEPT")
     os.system ("iptables -F")
     os.system ("iptables -X")
     os.system ("iptables -A INPUT -p tcp -i eth0 --dport 22 -j ACCEPT")
-    append_remove_port('22','TCP','a',False)
+    update_yml('22','TCP','a',False)
 
 
     os.system ("ip6tables -P INPUT ACCEPT")
     os.system ("ip6tables -F")
     os.system ("ip6tables -X")
     os.system ("ip6tables -A INPUT -p tcp -i eth0 --dport 22 -j ACCEPT")
-    append_remove_port('22','TCP','a',True)
+    update_yml('22','TCP','a',True)
 
-    for i,port in enumerate (TCP_port_list_ipv4):
+    for i,port in enumerate (firewall['ipv4']['TCP']):
         os.system ("iptables -A INPUT -p tcp -i eth0 --dport "+ str(port) +" -j ACCEPT")
         print("Port "+str(port)+" on protocol TCP  with ipv4 Open")
 
-    for i,port in enumerate (UDP_port_list_ipv4):
+    for i,port in enumerate (firewall['ipv4']['UDP']):
         os.system ("iptables -A INPUT -p udp -i eth0 --dport "+ str(port) +" -j ACCEPT")
         print("Port "+str(port)+" on protocol UDP  with ipv4 Open")
 
-    for i,port in enumerate (TCP_port_list_ipv6):
+    for i,port in enumerate (firewall['ipv6']['TCP']):
         os.system ("ip6tables -A INPUT -p tcp -i eth0 --dport "+ str(port) +" -j ACCEPT")
         print("Port "+str(port)+" on protocol TCP  with ipv6 Open")
 
-    for i,port in enumerate (UDP_port_list_ipv6):
+    for i,port in enumerate (firewall['ipv6']['UDP']):
         os.system ("ip6tables -A INPUT -p udp -i eth0 --dport "+ str(port) +" -j ACCEPT")
         print("Port "+str(port)+" on protocol UDP  with ipv6 Open")
 
@@ -140,14 +117,10 @@ def firewall_reload():
 
 
 
-def append_remove_port(port=None,protocol=None,mode=None,ip=None):
+def update_yml(port=None,protocol=None,mode=None,ip=None):
     '''
     Append port in firewall.yml
     '''
-    if ip == True:
-        ip = 'ipv6'
-    else:
-        ip = 'ipv4'
     
     with open('firewall.yml','r') as f:
         firewall = yaml.load(f)
@@ -164,12 +137,11 @@ def append_remove_port(port=None,protocol=None,mode=None,ip=None):
         else:
             print("Port already close")
     firewall[ip][protocol].sort()
-    f.close
 
     os.system("mv firewall.yml firewall.yml.old")
     with open('firewall.yml','w') as f:
         yaml.dump(firewall,f)
-        f.close
+