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