def getVulnByID(id, table_name, env):
"""
:param id: VulnerabilitiesObject id
:return: if id exists - returns VulnerabilitiesObject that is described by that id in the DB
"""
with sqlite3.connect(config.get('VulnServiceDB', env)) as db:
cursor = db.cursor()
cursor.execute("""SELECT * from %s where id = '%s'""" %
(table_name, id))
item = cursor.fetchone()
if (item is None):
raise Exception("No such vulnerability with id %s" % id)
if item[1] == config.get('SQLITypes', 'second_order'):
return SimpleVulnerabilityEntity(id=item[0],
name=item[1],
url=item[2],
payload=item[3],
requestB64=item[4],
affected_urls=json.loads(item[5]))
else:
return SimpleVulnerabilityEntity(id=item[0],
name=item[1],
url=item[2],
payload=item[3],
requestB64=item[4])
def getVulns(env, table_name, size=10, page=0):
"""
:param size: page size, default 10
:param page: page number, default 0
:return: a list of VulnerabilitiesObjects items from Vulns_Objects DB from page #page of size size
"""
with sqlite3.connect(config.get('VulnServiceDB', env)) as db:
cursor = db.cursor()
cursor.execute(
"""SELECT * from %s ORDER BY id ASC LIMIT %d OFFSET %d""" %
(table_name, size, page * size))
vulns_list = []
for vuln in cursor.fetchall():
if vuln[1] == config.get('SQLITypes', 'second_order'):
vulnEntity = SimpleVulnerabilityEntity(
id=vuln[0],
name=vuln[1],
url=vuln[2],
payload=vuln[3],
requestB64=vuln[4],
affected_urls=json.loads(vuln[5]))
else:
vulnEntity = SimpleVulnerabilityEntity(id=vuln[0],
name=vuln[1],
url=vuln[2],
payload=vuln[3],
requestB64=vuln[4])
vulns_list.append(vulnEntity)
return vulns_list
def setUp(self):
vuln_description1 = VulnerabilityDescriptionEntity(name='error-based', severity=1, description='abc',
recommendations='aaa')
vuln_description2 = VulnerabilityDescriptionEntity(name='RXSS', severity=2, description='def',
recommendations='bbb')
self.__vulnDescriptor.createVulnerabilityDescription(vuln_description1, self.env)
self.__vulnDescriptor.createVulnerabilityDescription(vuln_description2, self.env)
self.vuln1 = SimpleVulnerabilityEntity(name='error-based', url='http://www.something.com', payload='abcTest',
requestB64='aa+=')
self.vuln2 = SimpleVulnerabilityEntity(name='RXSS', url='http://www.anothersomething.com', payload='defTest',
requestB64='bb==')
self.vuln1ID = self.__VulnCrud.createVulnerability(self.vuln1, self.table_name, self.env).getID()
self.vuln2ID = self.__VulnCrud.createVulnerability(self.vuln2, self.table_name, self.env).getID()
def setUpClass(cls):
cls.vulnEntity = SimpleVulnerabilityEntity(id=1, name="rxss", url="http://test.com",
payload="<script>alert(1)<script>", requestB64="YmFzZTY0")
cls.vulnDesciption = VulnerabilityDescriptionEntity(name="DDOS", severity="10",
description="DDOS attack can harm the availablity of the application",
recommendations="Implement Carbon Black EDR")
cls.vulnBoundry = VulnerabilityBoundary(cls.vulnEntity, cls.vulnDesciption)
def convertToSimpleVulnerabilityEntity(self):
vulnEntity = SimpleVulnerabilityEntity(
id=self.__vulnID,
name=self.__name,
url=self.__url,
affected_urls=self.__affected_urls,
payload=self.__payload,
requestB64=self.__requestB64)
return vulnEntity
def add_event(self, name=None, url=None, payload=None, requestB64=None, affected_urls=None):
simpleVulnerability = SimpleVulnerabilityEntity(name=name, url=url, payload=payload, requestB64=requestB64, affected_urls=affected_urls)
createdVuln = VulnerabilitiesCRUD.createVulnerability(simpleVulnerability, self.__tableName, self.env_type)
print(createdVuln.getRequestB64())
def test_update_wrong_id(self):
with self.assertRaises(Exception) as cm:
self.__VulnCrud.updateVuln(
SimpleVulnerabilityEntity(id=self.vuln2ID + self.vuln1ID, name=self.vuln2.getName(),
url='http://www.something.com', payload='testUpdate', requestB64='aa+='),
self.table_name, self.env)
def test_update(self):
self.__VulnCrud.updateVuln(SimpleVulnerabilityEntity(id=self.vuln2ID, name=self.vuln2.getName(),
url='http://www.something.com', payload='testUpdate',
requestB64='aa+='), self.table_name, self.env)
self.assertEqual('testUpdate', self.__VulnCrud.getVulnByID(self.vuln2ID, self.table_name, self.env).getPayload())
def test_wrong_create_vulnerability(self):
self.assertNotEqual('abdTest', self.__VulnCrud.getVulns(self.env, self.table_name, 1, 0)[0].getPayload())
with self.assertRaises(Exception):
self.__VulnCrud.createVulnerability(
SimpleVulnerabilityEntity(name='a', url='http://www.something.com', payload='abcTest',
requestB64='aa+='), self.env)
class TestVulnerabilitiesCRUD(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.env = "prod"
cls.__VulnCrud = VulnerabilitiesCRUD
cls.__vulnDescriptor = VulnerabilityDescriptionCRUD
cls.table_name = 'vulns' + str(datetime.now()).replace('-', '').replace(' ', '').replace(':', '').replace('.', '')
cls.__VulnCrud.createTable(cls.table_name, cls.env)
VulnerabilityDescriptionCRUD.createTable(cls.env)
@classmethod
def tearDownClass(cls):
cls.__VulnCrud.dropTable(cls.table_name, cls.env)
cls.__VulnCrud = None
cls.__vulnDescriptor = None
def setUp(self):
vuln_description1 = VulnerabilityDescriptionEntity(name='error-based', severity=1, description='abc',
recommendations='aaa')
vuln_description2 = VulnerabilityDescriptionEntity(name='RXSS', severity=2, description='def',
recommendations='bbb')
self.__vulnDescriptor.createVulnerabilityDescription(vuln_description1, self.env)
self.__vulnDescriptor.createVulnerabilityDescription(vuln_description2, self.env)
self.vuln1 = SimpleVulnerabilityEntity(name='error-based', url='http://www.something.com', payload='abcTest',
requestB64='aa+=')
self.vuln2 = SimpleVulnerabilityEntity(name='RXSS', url='http://www.anothersomething.com', payload='defTest',
requestB64='bb==')
self.vuln1ID = self.__VulnCrud.createVulnerability(self.vuln1, self.table_name, self.env).getID()
self.vuln2ID = self.__VulnCrud.createVulnerability(self.vuln2, self.table_name, self.env).getID()
def tearDown(self):
self.__VulnCrud.deleteAllDataFromTable(self.table_name, self.env)
self.__vulnDescriptor.deleteAllDataFromTable(self.env)
def test_create_vulnerability(self):
self.assertEqual(self.vuln1.getPayload(), self.__VulnCrud.getVulns(self.env, self.table_name, 1, 0)[0].getPayload())
self.assertEqual(self.vuln2.getPayload(), self.__VulnCrud.getVulns(self.env, self.table_name, 1, 1)[0].getPayload())
def test_wrong_create_vulnerability(self):
self.assertNotEqual('abdTest', self.__VulnCrud.getVulns(self.env, self.table_name, 1, 0)[0].getPayload())
with self.assertRaises(Exception):
self.__VulnCrud.createVulnerability(
SimpleVulnerabilityEntity(name='a', url='http://www.something.com', payload='abcTest',
requestB64='aa+='), self.env)
def test_get_vulnerabilities_pagination(self):
self.assertEqual(len(self.__VulnCrud.getVulns(self.env, self.table_name, 2, 0)), 2)
def test_read_by_id(self):
self.assertEqual(self.vuln1.getPayload(), self.__VulnCrud.getVulnByID(self.vuln1ID, self.table_name, self.env).getPayload())
self.assertEqual(self.vuln2.getPayload(), self.__VulnCrud.getVulnByID(self.vuln2ID, self.table_name, self.env).getPayload())
def test_wrong_read_by_id(self):
with self.assertRaises(Exception):
self.__VulnCrud.getVulnByID(self.vuln2ID + self.vuln1ID, self.table_name, self.env)
def test_create_correct_number_of_vulnerabilities(self):
self.assertEqual(2, len(self.__VulnCrud.getVulns(self.env, self.table_name)))
def test_update(self):
self.__VulnCrud.updateVuln(SimpleVulnerabilityEntity(id=self.vuln2ID, name=self.vuln2.getName(),
url='http://www.something.com', payload='testUpdate',
requestB64='aa+='), self.table_name, self.env)
self.assertEqual('testUpdate', self.__VulnCrud.getVulnByID(self.vuln2ID, self.table_name, self.env).getPayload())
def test_update_wrong_id(self):
with self.assertRaises(Exception) as cm:
self.__VulnCrud.updateVuln(
SimpleVulnerabilityEntity(id=self.vuln2ID + self.vuln1ID, name=self.vuln2.getName(),
url='http://www.something.com', payload='testUpdate', requestB64='aa+='),
self.table_name, self.env)
def test_delete_by_id(self):
self.__VulnCrud.deleteVulnByID(self.vuln1ID, self.table_name, self.env)
self.assertEqual(1, len(self.__VulnCrud.getVulns(self.env, self.table_name)))
def test_delete_all_data_from_table(self):
self.__VulnCrud.deleteAllDataFromTable(self.table_name, self.env)
self.assertEqual(0, len(self.__VulnCrud.getVulns(self.env, self.table_name)))
def doCleanups(self):
pass
def suite(self):
pass