now faraday only uses the new models

model/common.py

@@ -550,40 +550,39 @@ def listModelObjectTypes(self):
         names.sort()
         return names
 
-    def generateID(self, classname, parent_id=None, **objargs):
-        if classname == 'Host':
-            args = (objargs['name'],)
-        elif classname == 'Interface':
-            args =(objargs['network_segment'], objargs['ipv4_address'], objargs['ipv6_address'])
-        elif classname == "Service":
-            args = (objargs['protocol'], ':'.join(str(objargs['ports'])))
-        elif classname == "Note":
-            args = (objargs['name'], objargs['text'])
-        elif classname == "Vulnerability":
-            args = (objargs['name'], objargs['desc'])
-        elif classname == "VulnerabilityWeb":
-            args = (objargs['name'], objargs['website'])
-        elif classname == "Credential" or classname == 'Cred':
-            print objargs
-            args = (objargs['name'], objargs['password'])
-        else:
-            raise Exception("You've provided the invalid classname {0}".format(classname))
-
-        objid = get_hash(args)
-        if parent_id:
-            objid = '.'.join([parent_id, objid])
-        return objid
+    def generateID(self, classname, parent_id='', **objargs):
+        # see how nicely formated that dictionary is
+        # it's a building about to go down on an eathquake!
+        # let's try not to make that an analogy about my code, ok? thank you :)
+        appropiate_class = self._registered_objects[classname]
+        class_to_args = {'Host': (objargs.get('name'),),
+                         'Cred': (objargs.get('name'), objargs.get('password')),
+                         'Note': (objargs.get('name'),
+                                  objargs.get('text')),
+                         'Service': (objargs.get('protocol'),
+                                     objargs.get('ports')),
+                         'Interface': (objargs.get('network_segment'),
+                                       objargs.get('ipv4_address'),
+                                       objargs.get('ipv6_address')),
+                         'Vulnerability': (objargs.get('name'),
+                                           objargs.get('desc')),
+                         'VulnerabilityWeb': (objargs.get('name'),
+                                              objargs.get('website'))
+                         }
+        try:
+            id = appropiate_class.generateID(parent_id, *class_to_args[classname])
+        except KeyError:
+            raise Exception("You've provided an invalid classname")
+        return id
 
-    def createModelObject(self, classname, object_name=None, workspace_name=None, **objargs):
+    def createModelObject(self, classname, object_name, workspace_name=None, parent_id=None, **objargs):
         if not workspace_name:
             workspace_name = CONF.getLastWorkspace()
         if classname in self._registered_objects:
             if object_name is not None:
-                print "OBJECT NAME IS NOT NONE"
                 objargs['name'] = object_name
-                objargs['id'] = self.generateID(classname, **objargs)
                 objargs['_id'] = -1 # they still don't have a server id
-                print objargs
+                objargs['id'] = self.generateID(classname, parent_id, **objargs)
                 tmpObj = self._registered_objects[classname](objargs, workspace_name)
                 return tmpObj
             else:

persistence/server/models.py

@@ -14,6 +14,7 @@
 from threading import Lock
 from persistence.server import server
 from persistence.server.utils import (force_unique,
+                                      get_hash,
                                       get_host_properties,
                                       get_interface_properties,
                                       get_service_properties,
@@ -50,15 +51,13 @@ def func_wrapper(*args, **kwargs):
     return func_wrapper
 
 def _flatten_dictionary(dictionary):
-    print "ORIGINAL DICTIONARY: ", dictionary
     flattened_dict = {}
     if dictionary.get('_id'):
         flattened_dict['_id'] = dictionary['_id']
     if dictionary.get('id'):
         flattened_dict['id'] = dictionary['id']
     for k, v in dictionary.get('value', {}).items():
         flattened_dict[k] = v
-    print "FLATTENED DICTIONARY: ", flattened_dict
     return flattened_dict
 
 def _get_faraday_ready_objects(workspace_name, faraday_ready_object_dictionaries,
@@ -502,18 +501,23 @@ def test_server_url(url_to_test):
 
 class ModelBase(object):
     def __init__(self, obj, workspace_name):
-        print "WHAT IS OBJECT: "
-        print obj
         self._workspace_name = workspace_name
         self._server_id = obj.get('_id')
-        self.id = obj.get('id')
+        self.id = obj['id']
         self.name = obj.get('name')
         self.description = obj.get('description', "")
         self.owned = obj.get('owned')
         self.owner = obj.get('owner')
         self._metadata = obj.get('metadata', Metadata(self.owner))
         self.updates = []
 
+    @staticmethod
+    def generateID(parent_id, *args):
+        objid = get_hash(args)
+        if parent_id:
+            objid = '.'.join([parent_id, objid])
+        return objid
+
     @staticmethod
     def publicattrsrefs():
         return {'Description': 'description',
@@ -585,11 +589,15 @@ class Host(ModelBase):
 
     def __init__(self, host, workspace_name):
         ModelBase.__init__(self, host, workspace_name)
-        print host
         self.default_gateway = host.get('default_gateway')
         self.os = host.get('os', 'unkown')
         self.vuln_amount = int(host.get('vulns', 0))
 
+    @staticmethod
+    def generateID(_, name):
+        # empty arg so as to share same interface as other classes' generateID
+        return ModelBase.generateID('', name)
+
     @staticmethod
     def publicattrsrefs():
         publicattrs = dict(ModelBase.publicattrsrefs(), **{
@@ -655,6 +663,10 @@ def __init__(self, interface, workspace_name):
         self.amount_ports_closed   = 0
         self.amount_ports_filtered = 0
 
+    @staticmethod
+    def generateID(parent_id, network_segment, ipv4_address, ipv6_address):
+        return ModelBase.generateID(parent_id, network_segment, ipv4_address, ipv6_address)
+
     @staticmethod
     def publicattrsrefs():
         publicattrs = dict(ModelBase.publicattrsrefs(), **{
@@ -752,6 +764,11 @@ def __init__(self, service, workspace_name):
         self.status = service['status']
         self.vuln_amount = int(service.get('vulns', 0))
 
+    @staticmethod
+    def generateID(parent_id, protocol, ports):
+        ports = ':'.join(str(ports))
+        return ModelBase.generateID(parent_id, protocol, ports)
+
     @staticmethod
     def publicattrsrefs():
         publicattrs = dict(ModelBase.publicattrsrefs(), **{
@@ -799,13 +816,21 @@ class Vuln(ModelBase):
 
     def __init__(self, vuln, workspace_name):
         ModelBase.__init__(self, vuln, workspace_name)
+
+        # this next two lines are stupid but so is life so you should get used to it :)
+        self.description = vuln['desc']
         self.desc = vuln['desc']
+
         self.data = vuln.get('data')
         self.severity = vuln['severity']
         self.refs = vuln.get('refs')
         self.confirmed = vuln.get('confirmed', False)
         self.resolution = vuln.get('resolution')
 
+    @staticmethod
+    def generateID(parent_id, name, description):
+        return ModelBase.generateID(parent_id, name, description)
+
     @staticmethod
     def publicattrsrefs():
         publicattrs = dict(ModelBase.publicattrsrefs(), **{
@@ -906,6 +931,10 @@ def __init__(self, vuln_web, workspace_name):
         self.target = vuln_web.get('target')
         self.parent = vuln_web.get('parent')
 
+    @staticmethod
+    def generateID(parent_id, name, website):
+        return ModelBase.generateID(parent_id, name, website)
+
     @staticmethod
     def publicattrsrefs():
         publicattrs = dict(ModelBase.publicattrsrefs(), **{
@@ -1003,6 +1032,10 @@ def __init__(self, note, workspace_name):
         ModelBase.__init__(self, note, workspace_name)
         self.text = note['text']
 
+    @staticmethod
+    def generateID(parent_id, name, text):
+        return ModelBase.generateID(parent_id, name, text)
+
     def updateAttributes(self, name=None, text=None):
         if name is not None:
             self.name = name
@@ -1025,6 +1058,10 @@ def __init__(self, credential, workspace_name):
 
         self.password = credential['password']
 
+    @staticmethod
+    def generateID(parent_id, name, password):
+        return ModelBase.generateID(parent_id, name, password)
+
     def updateAttributes(self, username=None, password=None):
         if username is not None:
             self.username =username

persistence/server/utils.py

@@ -6,6 +6,7 @@
 See the file 'doc/LICENSE' for the license information
 
 '''
+import hashlib
 from persistence.server.server_io_exceptions import MoreThanOneObjectFoundByID
 
 def force_unique(lst):
@@ -20,6 +21,9 @@ def force_unique(lst):
     else:
         raise MoreThanOneObjectFoundByID(lst)
 
+def get_hash(parts):
+    return hashlib.sha1("._.".join(parts)).hexdigest()
+
 def get_object_properties(obj):
     return {'id': obj.getID(),
             'name': obj.getName(),