From 028e68c63288f864e8330d80f69827432ea1577c Mon Sep 17 00:00:00 2001
From: Rafay Ghafoor <rafayghafoor@pm.me>
Date: Sun, 27 May 2018 23:40:11 +0500
Subject: [PATCH] refactoring

---
 router.py | 68 ++++++++++++++++++++++++++++---------------------------
 utils.py  |  8 +++++++
 2 files changed, 43 insertions(+), 33 deletions(-)

diff --git a/router.py b/router.py
index 61fa72a..9d408f9 100755
--- a/router.py
+++ b/router.py
@@ -2,7 +2,7 @@
 A PTCL router class which allows basic functionality for PTCL router.
 
 Usage Example:
-# router is used as an instance for the Router class in all examples.
+# papi is used as an instance for the Router class in all examples.
 >>> from router import Router
 >>> papi = Router(gateway='192.168.1.1')      # Launches session for interacting with router
 >>>
@@ -18,12 +18,11 @@
 {'my-computer': ['macxx', '192.168.10.1', '23 Hours, 59 Minutes']}
 '''
 import re
-import sys
 
 import requests
 import bs4
 
-from utils import convert_time
+import utils
 
 class Router():
     '''
@@ -36,9 +35,25 @@ class Router():
     '''
     mac_pattern = re.compile(u'^([0-9A-Fa-f]{2}[:-]){5}([0-9A-Fa-f]{2})$')
 
+    week_days = {
+        "Mon": 1,
+        "Tue": 2,
+        "Wed": 4,
+        "Thu": 8,
+        "Fri": 16,
+        "Sat": 32,
+        "Sun": 64,
+        "Everyday": 127
+        }
+
 
     def __init__(self, gateway="192.168.1.1", username="admin", password="ptcl"):
-        self.gateway = "http://" + gateway + '/'
+        if utils.validate_gateway(gateway):
+            self.gateway = "http://" + gateway + '/'
+        else:
+            raise ValueError("Invalid Gateway format specified, Valid format is",
+                             "192.168.1.1")
+
         self.username = username
         self.password = password
         self.dev_info = {}      # Devices info
@@ -67,8 +82,7 @@ def scrape_page(self, url, params='', soup='n'):
                 return request_url
 
         except requests.exceptions.ConnectionError:
-            print("Internet Connection Down.\nExiting...")
-            sys.exit()
+            raise Exception("Internet Connection Down.\nExiting...")
 
 
     def get_session_key(self):
@@ -84,7 +98,7 @@ def get_session_key(self):
         return self.sessionKey
 
 
-    def dhcpinfo(self):
+    def dhcp(self):
         '''
         Gets information from dhcp page.
         Format:
@@ -119,7 +133,7 @@ def dhcpinfo(self):
         return self.dev_info
 
 
-    def stationinfo(self):
+    def station(self):
         '''
         Gets information about the connected devices.
         '''
@@ -133,8 +147,9 @@ def stationinfo(self):
 
         return self.active_dev
 
+
     #TODO if already username defined, raise Error
-    def set_time_limit(self, username="User_1", mac="", days="Everyday", start="1", end="24"):
+    def time_limit(self, username="User_1", mac="", days="Everyday", start="1", end="24"):
         '''
         Restricts user from using internet for limited time.
         Creates a user profile containing mac, days, start_time, end_time.
@@ -159,24 +174,12 @@ def set_time_limit(self, username="User_1", mac="", days="Everyday", start="1",
         # Tuesday is   2^1
         # Wednesday is 2^2
         
-        week_days = {
-        "Mon": 1,
-        "Tue": 2,
-        "Wed": 4,
-        "Thu": 8,
-        "Fri": 16,
-        "Sat": 32,
-        "Sun": 64,
-        "Everyday": 127
-        }
-
         gen_params = lambda days: {
-        'username': username,
-        'days': days, 'start_time': start,
+        'username': username, 'days': days, 'start_time': start,
         'end_time': end, 'sessionKey': self.get_session_key()
         }
 
-        start, end = convert_time(start_time=start, end_time=end)
+        start, end = utils.convert_time(start_time=start, end_time=end)
         days = days.split('-')
 
         if days and len(days) < 3:
@@ -194,9 +197,7 @@ def set_time_limit(self, username="User_1", mac="", days="Everyday", start="1",
 
         return 'Successful'
 		
-		
-
-
+    
     def web_filter(self, url):
         '''
         Block website temporarily/permanently (i.e Temporarily, when time is specified).
@@ -212,9 +213,8 @@ def block(self, mac):
         Example:
         >>> router.block('xx:xx:xx:xx:xx:xx')
         '''
-        self.session.get(self.gateway + "wlmacflt.cmd?action=add&rmLst={}&sessionKey={}".format(devmac, self.get_session_key()))
-        return 'Successful'
-
+        return self.session.get(self.gateway + "wlmacflt.cmd?action=add&rmLst={}&sessionKey={}".format(devmac, self.get_session_key())) 
+        
 
     def unblock(self, mac):
         '''
@@ -224,13 +224,15 @@ def unblock(self, mac):
         Example:
         >>> router.unblock('xx:xx:xx:xx:xx:xx')
         '''
-        self.session.get(self.gateway + "wlmacflt.cmd?action=remove&rmLst={}&sessionKey={}".format(udevmac, self.get_session_key()))
-        return 'Successful'
+        return self.session.get(self.gateway + "wlmacflt.cmd?action=remove&rmLst={}&sessionKey={}".format(udevmac, self.get_session_key()))
 
 
     def reboot(self):
         '''
         Reboots Router.
         '''
-        self.session.get(self.gateway + "rebootinfo.cgi?sessionKey={}".format(self.get_session_key()))
-        return 'Successful'
\ No newline at end of file
+        return self.session.get(self.gateway + "rebootinfo.cgi?sessionKey={}".format(self.get_session_key()))
+
+
+    def __repr__(self):
+        return self.gateway
\ No newline at end of file
diff --git a/utils.py b/utils.py
index 69647aa..b2fbaee 100644
--- a/utils.py
+++ b/utils.py
@@ -1,3 +1,11 @@
+import re
+
+def validate_gateway(gateway):
+    if not re.search("https?://", gateway) and not gateway.endswith('/'):
+        return True
+    return False
+
+
 def convert_time(start_time="1", end_time="23:59"):
         # TODO : Add test that the numbers after : shouldn't exceed 60 (minutes)
         '''