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