diff options
Diffstat (limited to 'lib/model.py')
-rw-r--r-- | lib/model.py | 181 |
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 |