class Test_HostsModel(Common):
"""Unit tests for HostsModel"""
model = None
def setup(self):
self.model = HostsModel()
self.db.q("TRUNCATE TABLE `hosts`")
self.db.q("TRUNCATE TABLE `ips`")
def test_get_id_by_name(self):
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (2,1,1,'test.com', '')")
assert self.model.get_id_by_name(1, 'test.com') == 2
def test_exists(self):
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (2,1,1,'test.com', '')")
assert self.model.exists(1, 'test.com')
assert not self.model.exists(1, 'test1.com')
def test_add(self):
assert not self.model.exists(1, 'test.com')
assert self.model.add(1, 1, 'test.com', '') == 1
assert self.model.exists(1, 'test.com')
def test_list(self):
self.db.q("INSERT INTO `ips` (id, project_id, ip, descr) VALUES (1, 1, '111.111.111.111', '')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (1, 1, 1, 'test1.com', 'desc1')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (2, 1, 1, 'test2.com', 'desc2')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (3, 1, 1, 'test3.com', 'desc3')")
test_data = [
{'id': 1, 'name': 'test1.com', 'ip': '111.111.111.111', 'descr': 'desc1'},
{'id': 2, 'name': 'test2.com', 'ip': '111.111.111.111', 'descr': 'desc2'},
{'id': 3, 'name': 'test3.com', 'ip': '111.111.111.111', 'descr': 'desc3'},
]
assert test_data == self.model.list(1, '111.111.111.111')
def test_list_of_names(self):
self.db.q("INSERT INTO `ips` (id, project_id, ip, descr) VALUES (1, 1, '111.111.111.111', '')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (1, 1, 1, 'test1.com', 'desc1')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (2, 1, 1, 'test2.com', 'desc2')")
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (3, 1, 1, 'test3.com', 'desc3')")
test_data = ['test1.com', 'test2.com', 'test3.com']
assert test_data == self.model.list_of_names(1)
def test_delete(self):
self.db.q("INSERT INTO `hosts` (id, project_id, ip_id, name, descr) VALUES (2,1,1,'test.com', '')")
assert self.model.exists(1, 'test.com')
self.model.delete(1, 'test.com')
assert not self.model.exists(1, 'test.com')
def brute_action(self):
""" Action brute of module """
self.enable_logger()
self.validate_main()
self.pre_start_inf()
q = DnsBruteJob()
loaded = self.load_objects(q)
self.logger.log(
"Loaded {0} words ({1}-{2}) from all {3}.".format(
(loaded['end'] - loaded['start']), loaded['start'], loaded['end'], loaded['all'])
if (int(self.options['parts'].value) and int(self.options['part'].value)) else
"Loaded {0} words from source.".format(loaded['all'])
)
counter = WSCounter(5, 300, loaded['all'] if not loaded['end'] else loaded['end']-loaded['start'])
result = []
w_thrds = []
DnsRoller = Roller()
DnsRoller.load_file(Registry().get('wr_path') + '/bases/dns-servers.txt')
for _ in range(int(self.options['threads'].value)):
we_need_server = True
while we_need_server:
we_need_server = False
try:
next_server = DnsRoller.get()
#print "Next DNS " + next_server
if self.options['protocol'].value == 'auto':
try:
dns.query.tcp(dns.message.make_query('test.com', 'A'), next_server, timeout=5)
protocol = 'tcp'
except socket.error:
try:
dns.query.udp(dns.message.make_query('test.com', 'A'), next_server, timeout=5)
protocol = 'udp'
except socket.error:
#raise Exception('Can`t detect DNS-server protocol. Check addr.')
we_need_server = True
#print 'DNS protolol detected automaticaly: ' + protocol
else:
protocol = self.options['protocol'].value
except dns.exception.Timeout:
self.logger.log("Check server {0}. Don`t work.".format(next_server))
we_need_server = True
hosts = []
if self.options['host'].value == 'all':
hosts_model = HostsModel()
hosts.extend(hosts_model.list_of_names(Registry().get('pData')['id']))
else:
hosts.append(self.options['host'].value)
worker = DnsBruteThread(
q,
hosts,
self.options['template'].value,
protocol,
self.options['msymbol'].value,
self.options['ignore-ip'].value,
next_server,
self.options['delay'].value,
self.options['http-not-found-re'].value,
self.options['ignore-words-re'].value,
result,
counter
)
worker.setDaemon(True)
worker.start()
w_thrds.append(worker)
time.sleep(1)
while len(w_thrds):
for worker in w_thrds:
if worker.done or Registry().get('positive_limit_stop'):
del w_thrds[w_thrds.index(worker)]
time.sleep(2)
if self.options['host'].value == 'all':
self._output_zones(result)
else:
self._output(result)
self.done = True