Merge pull request #24 from bwall/dev

Updated zoption to handle multiple data types
hatRiot · Jun 23, 2015 · 02eebd4 · 02eebd4
2 parents 59986c6 + b7e3bfb
commit 02eebd4
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 18 deletions.
diff --git a/src/core/util.py b/src/core/util.py
@@ -438,12 +438,6 @@ def eval_type(value, type):
     elif type == "str":
         # anything can be a string
         rval = (True, str(value))
-    elif type == "ip or ipmask":
-        t1 = eval_type(value, "ip")
-        if t1[0] == False:
-            t1 = eval_type(value, "ipmask")
-            return t1
-        return t1
     elif type == "ipmask":
         ip = value.split('.')
         if len(ip) != 4:

diff --git a/src/core/zoption.py b/src/core/zoption.py
@@ -1,16 +1,22 @@
 from util import eval_type
 
+
 class Zoption:
     """ generic option class for managing and validating
         zarp options.
     """
-    def __init__(self, value = None, type = None, required = False, 
-                       display = None, opts = None):
+
+    def __init__(self, value=None, type=None, required=False, display=None, opts=None):
         self.value = value
-        self.type = type
+        if isinstance(type, basestring):
+            self.types = [type]
+            self.type = type
+        else:
+            self.types = type
+            self.type = None
         self.required = required
         self.display = display
-        self.opts    = opts
+        self.opts = opts
 
     def getStr(self):
         """ Some objects don't have a __str__ method (regex),
@@ -30,9 +36,10 @@ def validate(self):
         """ Validates the object's value to ensure it conforms
             to whatever type the object dictates.
         """
-        rvals = eval_type(self.value, self.type)
-        if rvals[0]:
-            self.value = rvals[1]
-            return True
-        else:
-            return False
+        for t in self.types:
+            rvals = eval_type(self.value, t)
+            if rvals[0]:
+                self.value = rvals[1]
+                self.type = t
+                return True
+        return False
diff --git a/src/modules/poison/arp.py b/src/modules/poison/arp.py
@@ -30,7 +30,7 @@ def __init__(self):
                                              required=True,
                                              display="Target to poison"),
                             "from_ip": Zoption(value=None,
-                                               type="ip or ipmask",
+                                               type=["ip", "ipmask"],
                                                required=False,
                                                display="Address or addresses to spoof from target"),
                             "respoof": Zoption(value=2,
@@ -254,4 +254,4 @@ def view(self):
         """ ARP poisoner doesnt have a view, yet.
         """
         Msg('No view for ARP poison.  Enable a sniffer for detailed analysis.')
-        return
+        return