def test_insert_functionality():
with DB('test.db') as db:
db.insert_rows(query=QUERIES['test_insert'].format('test_table'),
input_t=[[34]])
test_insert_output = list(
db.select_query(query=QUERIES['select_all'].format('test_table')))
assert test_insert_output == [(23, ), (34, )]
def test_table_manager():
with DB('test.db') as db:
db.table_manager(query=QUERIES['test_create'].format('test_table_2'))
assert list(
db.select_query(
query=QUERIES['exist'].format('test_table_2'))) == [
('test_table_2', )
]
db.table_manager(query=QUERIES['drop'].format('test_table_2'))
assert list(
db.select_query(
query=QUERIES['exist'].format('test_table_2'))) == []
db.table_manager('')
def lookup_vulnerabilities_in_database(product_name: str,
requested_version: str) -> list:
with DB(str(Path(__file__).parent.parent / 'internal' / DB_NAME)) as db:
product_terms, version = unbinding(
generate_search_terms(product_name)), unbinding(
[requested_version])[0]
matched_cpe = list(match_cpe(db, product_terms))
if len(matched_cpe) == 0:
print('No CPEs were found!\n')
return ['N/A']
else:
matched_product = sort_cpe_matches(matched_cpe, version)
cve_candidates = list(set(search_cve(db, matched_product)))
cve_candidates.extend(
list(set(search_cve_summary(db, matched_product))))
return cve_candidates
def setup() -> None:
try:
remove('cve_cpe.db')
except OSError:
pass
cpe_base = dp.setup_cpe_table(
dp.extract_cpe(PATH_TO_TEST + EXTRACT_CPE_XML))
cve_base, summary_base = dp.extract_cve(PATH_TO_TEST + EXTRACT_CVE_JSON)
cve_base, summary_base = dp.setup_cve_feeds_table(
cve_list=cve_base), dp.setup_cve_summary_table(
summary_list=summary_base)
with DB(PATH_TO_TEST + 'test_update.db') as db:
db.table_manager(query=QUERIES['create_cpe_table'].format('cpe_table'))
db.insert_rows(query=QUERIES['insert_cpe'].format('cpe_table'),
input_t=cpe_base)
db.table_manager(query=QUERIES['create_cve_table'].format('cve_table'))
db.table_manager(
query=QUERIES['create_summary_table'].format('summary_table'))
db.insert_rows(query=QUERIES['insert_cve'].format('cve_table'),
input_t=cve_base)
db.insert_rows(query=QUERIES['insert_summary'].format('summary_table'),
input_t=summary_base)
db.table_manager(
query=QUERIES['test_create_update'].format('outdated'))
db.table_manager(query=QUERIES['test_create_update'].format('new'))
db.insert_rows(query=QUERIES['test_insert_cve_id'].format('outdated'),
input_t=[('CVE-2018-0001', 2018),
('CVE-2018-0002', 2018)])
db.insert_rows(query=QUERIES['test_insert_cve_id'].format('new'),
input_t=[('CVE-2018-0002', 2018),
('CVE-2018-0003', 2018)])
yield None
try:
remove(PATH_TO_TEST + 'test_update.db')
remove(PATH_TO_TEST + 'test_import.db')
remove(PATH_TO_TEST + 'test_output.db')
except OSError:
pass
def test_select_functionality():
with DB('test.db') as db:
assert list(
db.select_query(
query=QUERIES['select_all'].format('test_table'))) == [(23, )]
def test_db_connection():
with DB('test.db') as db:
assert db.conn is not None
with pytest.raises(TypeError):
DB('')
def match_cpe(db: DB,
product_search_terms: list) -> Generator[namedtuple, None, None]:
for vendor, product, version in db.select_query(QUERIES['cpe_lookup']):
for product_term in product_search_terms:
if terms_match(product_term, product):
yield PRODUCT(vendor, product, version)
def search_cve_summary(db: DB,
product: namedtuple) -> Generator[str, None, None]:
for cve_id, summary in db.select_query(QUERIES['summary_lookup']):
if product_is_in_wordlist(product, summary.split(' ')):
yield cve_id
def search_cve(db: DB, product: namedtuple) -> Generator[str, None, None]:
for cve_id, vendor, product_name, version in db.select_query(
QUERIES['cve_lookup']):
if terms_match(product.vendor_name, vendor) and terms_match(product.product_name, product_name) \
and (product.version_number.startswith(get_version_index(version, 0)) or version == 'ANY' or version == 'NA'):
yield cve_id