Global refactoring

1 files changed, 181 insertions, 0 deletions

diff --git a/lib/model.py b/lib/model.py
new file mode 100644
index 0000000..933d901
--- /dev/null
+++ b/lib/model.py
@@ -0,0 +1,181 @@
+from xml.etree.ElementTree import ElementTree
+
+from lib.config import Config
+
+class Network():
+    """A simple network definition"""
+    
+    def __init__(self):
+        self.name = ""
+        self.domains = []
+        self.admins = []
+        self.ip_blocks = []
+        self.data = {}
+
+    def as_dict(self, data):
+        # Beautify
+        result = {}
+        domains = data.get_domains()
+        persons = data.get_persons()
+
+        result['name'] = self.name
+
+        # Domains
+        result['domains'] = []
+        for d in self.domains:
+            result['domains'].append(domains[d])
+
+        # Admins
+        result['admins'] = []
+        for a in self.admins:
+            result['admins'].append(persons[a]) 
+
+        # Networks
+        result['networks'] = self.ip_blocks
+
+        return result        
+
+class Domain():
+    """A simple domain definition"""
+
+    def __init__(self):
+        self.name = ""
+        self.admins = []
+        self.data = {}
+
+    def add_admin(self, admin):
+        """Add an administrator for this network"""
+
+        self.admins.append(admin)
+
+class Person():
+    """A simple person definition"""
+
+    def __init__(self):
+        self.name = ""
+        self.surname = ""
+        self.email = ""
+        self.data = {}
+
+class Data():
+    """Abstract class for storing and getting information"""
+    
+    def __init__(self):
+        self.config = Config().parser
+        self.networks = None
+        self.domains = None
+        self.persons = None
+
+    def parse_config(self):
+        """Parse neccesary config params depending on the method used
+
+        Abstract method"""
+        
+        pass
+    
+    def load_data(self):
+        """Load data from defined source.
+
+        Abstract method"""
+        
+        pass
+
+    def get_networks(self):
+        """Return all networks. Common method for all kind of storages."""
+        
+        if self.networks == None:
+            self.load_data()
+
+        return self.networks
+
+    def get_domains(self):
+        """Return all domains. Common method for all kind of storages."""
+        
+        if self.networks == None:
+            self.load_data()
+
+        return self.domains
+
+    def get_persons(self):
+        """Return all persons. Common method for all kind of storages."""
+        
+        if self.persons == None:
+            self.load_data()
+
+        return self.persons
+
+class DataXML(Data):
+    """Reads network information from a XML file"""
+    
+    def parse_config(self):
+        """Reads and sets up XML config file fields"""
+        
+        self.data_file = self.config['Storage']['xml_file']
+        
+    def load_data(self):
+        """Parse XML for getting network information"""
+
+        self.parse_config()
+        root = ElementTree(file=self.data_file).getroot()
+        
+        self.networks = {}
+        self.domains = {}
+        self.persons = {}
+        
+        for elem in root:
+            if elem.tag == 'person':
+                self.add_person(elem)
+            elif elem.tag == 'domain':
+                self.add_domain(elem)
+            elif elem.tag == 'network':
+                self.add_network(elem)
+            else:
+                pass # raise TagNotFoundError
+
+    def add_person(self, elem):
+        """Adds a new person"""
+
+        person = Person()
+        for e in elem:
+            if e.tag == 'name':
+                person.name = e.text
+            elif e.tag == 'surname':
+                person.surname = e.text
+            elif e.tag == 'email':
+                person.email = e.text
+            else:
+                person.data[e.tag] = e.text
+
+        print("[+] Read person: {0} - {1} - {2}".format(person.name, person.surname, person.email))
+        self.persons[elem.attrib['id']] = person
+
+    def add_domain(self, elem):
+        """Adds a new domain"""
+
+        domain = Domain()
+        for e in elem:
+            if e.tag == 'name':
+                domain.name = e.text
+            else:
+                domain.data[e.tag] = e.text
+
+        print("[+] Read domain: {}".format(domain.name))
+        self.domains[elem.attrib['id']] = domain                
+
+    def add_network(self, elem):
+        """Adds a new network"""
+        
+        network = Network()
+        for e in elem:
+            if e.tag == 'name':
+                network.name = e.text
+            elif e.tag == 'domain':
+                network.domains.append(e.text)            
+            elif e.tag == 'ip_block':
+                network.ip_blocks.append(e.text)
+            elif e.tag == 'admin':
+                network.admins.append(e.text)
+            else:
+                network.data[e.tag] = e.text
+                
+        self.networks[elem.attrib['id']] = network