Fixed arguments

RafayGhafoor · Sep 10, 2017 · 9d5c4f9 · 9d5c4f9
1 parent 93dd430
commit 9d5c4f9
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 56 deletions.
diff --git a/ptcl.py b/ptcl.py
@@ -76,73 +76,66 @@ def main():
         configure.config_check()
         sys.exit("Please Re-run.")
 
-    if args.cli == 'False':
-        my_macs = configure.get_alias()
-        if args.block:
-            # print "Calling blocker Function"
-            if args.block in my_macs.iterkeys():
-                # print "Calling blocker function - AUTOMATED MODE."
-                ptcl.block_dev(my_macs[args.block.lower()])
-                print "%s has been blocked." % args.block.capitalize()
-                if args.block not in my_macs.iterkeys():
-                    print "User not found."
-            elif args.block not in my_macs.iterkeys():
-                print "User not found."
-
-        elif args.unblock:
-            if args.unblock in my_macs.iterkeys():
-                # print "Calling unblocker function - AUTOMATED MODE"
-                ptcl.unblock_dev(my_macs[args.unblock.lower()])
-                print "%s has been unblocked." % args.unblock.capitalize()
-            elif args.unblock not in my_macs.iterkeys():
-                print "User not found."
+    my_macs = configure.get_alias()
 
-        elif args.active_devices:
-            # print "Calling Station info Function"
-            ptcl.get_stationinfo()
-            print "Currently active devices are:", len(ptcl.active_dev)
+    if args.block:
+        if args.block == '1':
+            name = show_active_dev()
+            ptcl.host_and_mac = dict(ptcl.host_and_mac)
+            dev_mac = int(raw_input("Please Enter Device Number: ")) - 1
+            ptcl.block_dev(ptcl.host_and_mac[name[dev_mac]])
+            print "%s has been blocked." % name[dev_mac].capitalize()
 
-        elif args.restart:
-            # print "Calling restart Function"
-            ptcl.reboot()
+        elif args.block != '1' and args.block in my_macs.iterkeys():
+            # print "Calling blocker function - AUTOMATED MODE."
+            ptcl.block_dev(my_macs[args.block.lower()])
+            print "%s has been blocked." % args.block.capitalize()
+            if args.block not in my_macs.iterkeys():
+                print "User not found."
 
-        elif args.show_dhcp:
-            # print "Calling DHCP_info Function"
-            show_dhcpinfo()
+        elif args.block != '1' and args.block not in my_macs.iterkeys():
+            print "User not found."
 
-        elif args.blocked_dev:
+    elif args.unblock:
+        if args.unblock == '1':
             show_blocked_dev()
+            dev_mac = raw_input("Please enter device mac address: ")
+            ptcl.unblock_dev(dev_mac)
+            print "%s has been unblocked." % dev_mac
 
-        elif args.set_alias:
-            show_active_dev()
-            configure.set_alias()
+        elif args.unblock != 1 and args.unblock in my_macs.iterkeys():
+            # print "Calling unblocker function - AUTOMATED MODE"
+            ptcl.unblock_dev(my_macs[args.unblock.lower()])
+            print "%s has been unblocked." % args.unblock.capitalize()
 
-        elif args.show_active == '.':
-            # print "Calling show_active Function"
-            show_active_dev()
+        elif args.unblock != 1 and args.unblock not in my_macs.iterkeys():
+            print "User not found."
 
-        else:
-            print "Invalid Argument"
+    elif args.active_devices:
+        # print "Calling Station info Function"
+        ptcl.get_stationinfo()
+        print "Currently active devices are:", len(ptcl.active_dev)
 
+    elif args.restart:
+        # print "Calling restart Function"
+        ptcl.reboot()
 
-    elif not args.cli:
-        if not args.block:
-            # print "Calling blocker function - CLI MODE."
-            name = show_active_dev()
-            ptcl.host_and_mac = dict(ptcl.host_and_mac)
-            dev_mac = int(raw_input("Please Enter Device Number: ")) - 1
-            ptcl.block_dev(ptcl.host_and_mac[name[dev_mac]])
-            print "%s has been blocked." % name[dev_mac].capitalize()
+    elif args.show_dhcp:
+        # print "Calling DHCP_info Function"
+        show_dhcpinfo()
 
+    elif args.blocked_dev:
+        show_blocked_dev()
 
-        elif not args.unblock:
-            # print "Calling unblocker function - CLI MODE."
-            name = show_active_dev()
-            ptcl.host_and_mac = dict(ptcl.host_and_mac)
-            dev_mac = raw_input("Please enter device mac address: ")
-            ptcl.unblock_dev(dev_mac)
-            # ptcl.unblock_dev(ptcl.host_and_mac[name[dev_mac]])
-            print "%s has been unblocked." % name[dev_mac].capitalize()
+    elif args.set_alias:
+        show_active_dev()
+        configure.set_alias()
 
+    elif args.show_active == '.':
+        # print "Calling show_active Function"
+        show_active_dev()
+
+    else:
+        print "Invalid Argument"
 
 main()
diff --git a/router.py b/router.py
@@ -7,7 +7,7 @@
 >>> from router import Router
 >>> router = Router('192.168.1.1')
 >>> router.reboot() # Reboots router
->>> router.active_dev() # Shows devices which are currently connected to the router
+>>> router.active_dev() # Returns a list of active devices.
 '''
 import requests
 import bs4
@@ -76,6 +76,7 @@ def dhcpinfo(self):
                 '''
                 self.dev_hostname.append(td[td.index(i) - 1].text.encode('ascii'))
                 self.mac_address.append(i.text.encode('ascii'))
+        return (self.dev_hostname, self.mac_address)
 
 
     def stationinfo(self):
@@ -86,6 +87,7 @@ def stationinfo(self):
         for i in soup.findAll('td'):
             if self.mac_adr_regex.search(i.text.strip()):
                 self.active_dev.append(i.text.strip().lower().encode('ascii'))
+        return self.active_dev
 
 
     def block_dev(self, devmac):